dictserver.py (6125B)
1 #!/usr/bin/env python3 2 # -*- coding: utf-8 -*- 3 #*************************************************************************** 4 # _ _ ____ _ 5 # Project ___| | | | _ \| | 6 # / __| | | | |_) | | 7 # | (__| |_| | _ <| |___ 8 # \___|\___/|_| \_\_____| 9 # 10 # Copyright (C) Daniel Stenberg, <daniel@haxx.se>, et al. 11 # 12 # This software is licensed as described in the file COPYING, which 13 # you should have received as part of this distribution. The terms 14 # are also available at https://curl.se/docs/copyright.html. 15 # 16 # You may opt to use, copy, modify, merge, publish, distribute and/or sell 17 # copies of the Software, and permit persons to whom the Software is 18 # furnished to do so, under the terms of the COPYING file. 19 # 20 # This software is distributed on an "AS IS" basis, WITHOUT WARRANTY OF ANY 21 # KIND, either express or implied. 22 # 23 # SPDX-License-Identifier: curl 24 # 25 ########################################################################### 26 # 27 """DICT server.""" 28 29 from __future__ import (absolute_import, division, print_function, 30 unicode_literals) 31 32 import argparse 33 import logging 34 import os 35 import sys 36 37 from util import ClosingFileHandler 38 39 try: # Python 2 40 import SocketServer as socketserver # type: ignore 41 except ImportError: # Python 3 42 import socketserver 43 44 log = logging.getLogger(__name__) 45 HOST = "localhost" 46 47 # The strings that indicate the test framework is checking our aliveness 48 VERIFIED_REQ = b"verifiedserver" 49 VERIFIED_RSP = "WE ROOLZ: {pid}" 50 51 52 def dictserver(options): 53 """Start up a TCP server with a DICT handler and serve DICT requests forever.""" 54 if options.pidfile: 55 pid = os.getpid() 56 # see tests/server/util.c function write_pidfile 57 if os.name == "nt": 58 pid += 4194304 59 with open(options.pidfile, "w") as f: 60 f.write(str(pid)) 61 62 local_bind = (options.host, options.port) 63 log.info("[DICT] Listening on %s", local_bind) 64 65 # Need to set the allow_reuse on the class, not on the instance. 66 socketserver.TCPServer.allow_reuse_address = True 67 server = socketserver.TCPServer(local_bind, DictHandler) 68 server.serve_forever() 69 70 return ScriptRC.SUCCESS 71 72 73 class DictHandler(socketserver.BaseRequestHandler): 74 """Handler class for DICT connections.""" 75 76 def handle(self): 77 """Respond to all queries with a 552.""" 78 try: 79 # First, send a response to allow the server to continue. 80 rsp = "220 dictserver <xnooptions> <msgid@msgid>\n" 81 self.request.sendall(rsp.encode("utf-8")) 82 83 # Receive the request. 84 data = self.request.recv(1024).strip() 85 log.debug("[DICT] Incoming data: %r", data) 86 87 if VERIFIED_REQ in data: 88 log.debug("[DICT] Received verification request from test " 89 "framework") 90 pid = os.getpid() 91 # see tests/server/util.c function write_pidfile 92 if os.name == "nt": 93 pid += 4194304 94 response_data = VERIFIED_RSP.format(pid=pid) 95 else: 96 log.debug("[DICT] Received normal request") 97 response_data = "No matches" 98 99 # Send back a failure to find. 100 response = "552 {0}\n".format(response_data) 101 log.debug("[DICT] Responding with %r", response) 102 self.request.sendall(response.encode("utf-8")) 103 104 except IOError: 105 log.exception("[DICT] IOError hit during request") 106 107 108 def get_options(): 109 parser = argparse.ArgumentParser() 110 111 parser.add_argument("--port", action="store", default=9016, 112 type=int, help="port to listen on") 113 parser.add_argument("--host", action="store", default=HOST, 114 help="host to listen on") 115 parser.add_argument("--verbose", action="store", type=int, default=0, 116 help="verbose output") 117 parser.add_argument("--pidfile", action="store", 118 help="file name for the PID") 119 parser.add_argument("--logfile", action="store", 120 help="file name for the log") 121 parser.add_argument("--srcdir", action="store", help="test directory") 122 parser.add_argument("--id", action="store", help="server ID") 123 parser.add_argument("--ipv4", action="store_true", default=0, 124 help="IPv4 flag") 125 126 return parser.parse_args() 127 128 129 def setup_logging(options): 130 """Set up logging from the command line options.""" 131 root_logger = logging.getLogger() 132 add_stdout = False 133 134 formatter = logging.Formatter("%(asctime)s %(levelname)-5.5s %(message)s") 135 136 # Write out to a logfile 137 if options.logfile: 138 handler = ClosingFileHandler(options.logfile) 139 handler.setFormatter(formatter) 140 handler.setLevel(logging.DEBUG) 141 root_logger.addHandler(handler) 142 else: 143 # The logfile wasn't specified. Add a stdout logger. 144 add_stdout = True 145 146 if options.verbose: 147 # Add a stdout logger as well in verbose mode 148 root_logger.setLevel(logging.DEBUG) 149 add_stdout = True 150 else: 151 root_logger.setLevel(logging.INFO) 152 153 if add_stdout: 154 stdout_handler = logging.StreamHandler(sys.stdout) 155 stdout_handler.setFormatter(formatter) 156 stdout_handler.setLevel(logging.DEBUG) 157 root_logger.addHandler(stdout_handler) 158 159 160 class ScriptRC(object): 161 """Enum for script return codes.""" 162 163 SUCCESS = 0 164 FAILURE = 1 165 EXCEPTION = 2 166 167 168 if __name__ == '__main__': 169 # Get the options from the user. 170 options = get_options() 171 172 # Setup logging using the user options 173 setup_logging(options) 174 175 # Run main script. 176 try: 177 rc = dictserver(options) 178 except Exception: 179 log.exception('Error running server') 180 rc = ScriptRC.EXCEPTION 181 182 if options.pidfile and os.path.isfile(options.pidfile): 183 os.unlink(options.pidfile) 184 185 log.info("[DICT] Returning %d", rc) 186 sys.exit(rc)