2 # -*- coding: utf-8 -*-
6 from __future__ import (absolute_import, division, print_function,
13 import SocketServer as socketserver
14 except ImportError: # Python 3
18 log = logging.getLogger(__name__)
21 # The strings that indicate the test framework is checking our aliveness
22 VERIFIED_REQ = b"verifiedserver"
23 VERIFIED_RSP = "WE ROOLZ: {pid}"
26 def dictserver(options):
28 Starts up a TCP server with a DICT handler and serves DICT requests
33 with open(options.pidfile, "w") as f:
34 f.write("{0}".format(pid))
36 local_bind = (HOST, options.port)
37 log.info("[DICT] Listening on %s", local_bind)
39 # Need to set the allow_reuse on the class, not on the instance.
40 socketserver.TCPServer.allow_reuse_address = True
41 server = socketserver.TCPServer(local_bind, DictHandler)
42 server.serve_forever()
44 return ScriptRC.SUCCESS
47 class DictHandler(socketserver.BaseRequestHandler):
48 """Handler class for DICT connections.
53 Simple function which responds to all queries with a 552.
56 # First, send a response to allow the server to continue.
57 rsp = "220 dictserver <xnooptions> <msgid@msgid>\n"
58 self.request.sendall(rsp.encode("utf-8"))
60 # Receive the request.
61 data = self.request.recv(1024).strip()
62 log.debug("[DICT] Incoming data: %r", data)
64 if VERIFIED_REQ in data:
65 log.debug("[DICT] Received verification request from test "
67 response_data = VERIFIED_RSP.format(pid=os.getpid())
69 log.debug("[DICT] Received normal request")
70 response_data = "No matches"
72 # Send back a failure to find.
73 response = "552 {0}\n".format(response_data)
74 log.debug("[DICT] Responding with %r", response)
75 self.request.sendall(response.encode("utf-8"))
78 log.exception("[DICT] IOError hit during request")
82 parser = argparse.ArgumentParser()
84 parser.add_argument("--port", action="store", default=9016,
85 type=int, help="port to listen on")
86 parser.add_argument("--verbose", action="store", type=int, default=0,
87 help="verbose output")
88 parser.add_argument("--pidfile", action="store",
89 help="file name for the PID")
90 parser.add_argument("--logfile", action="store",
91 help="file name for the log")
92 parser.add_argument("--srcdir", action="store", help="test directory")
93 parser.add_argument("--id", action="store", help="server ID")
94 parser.add_argument("--ipv4", action="store_true", default=0,
97 return parser.parse_args()
100 def setup_logging(options):
102 Set up logging from the command line options
104 root_logger = logging.getLogger()
107 formatter = logging.Formatter("%(asctime)s %(levelname)-5.5s %(message)s")
109 # Write out to a logfile
111 handler = logging.FileHandler(options.logfile, mode="w")
112 handler.setFormatter(formatter)
113 handler.setLevel(logging.DEBUG)
114 root_logger.addHandler(handler)
116 # The logfile wasn't specified. Add a stdout logger.
120 # Add a stdout logger as well in verbose mode
121 root_logger.setLevel(logging.DEBUG)
124 root_logger.setLevel(logging.INFO)
127 stdout_handler = logging.StreamHandler(sys.stdout)
128 stdout_handler.setFormatter(formatter)
129 stdout_handler.setLevel(logging.DEBUG)
130 root_logger.addHandler(stdout_handler)
133 class ScriptRC(object):
134 """Enum for script return codes"""
140 class ScriptException(Exception):
144 if __name__ == '__main__':
145 # Get the options from the user.
146 options = get_options()
148 # Setup logging using the user options
149 setup_logging(options)
153 rc = dictserver(options)
154 except Exception as e:
156 rc = ScriptRC.EXCEPTION
158 log.info("[DICT] Returning %d", rc)