• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1#!/usr/bin/env python
2# -*- coding: utf-8 -*-
3#
4""" DICT server """
5
6from __future__ import (absolute_import, division, print_function,
7                        unicode_literals)
8import argparse
9import os
10import sys
11import logging
12try:  # Python 2
13    import SocketServer as socketserver
14except ImportError:  # Python 3
15    import socketserver
16
17
18log = logging.getLogger(__name__)
19HOST = "localhost"
20
21# The strings that indicate the test framework is checking our aliveness
22VERIFIED_REQ = b"verifiedserver"
23VERIFIED_RSP = "WE ROOLZ: {pid}"
24
25
26def dictserver(options):
27    """
28    Starts up a TCP server with a DICT handler and serves DICT requests
29    forever.
30    """
31    if options.pidfile:
32        pid = os.getpid()
33        with open(options.pidfile, "w") as f:
34            f.write("{0}".format(pid))
35
36    local_bind = (options.host, options.port)
37    log.info("[DICT] Listening on %s", local_bind)
38
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()
43
44    return ScriptRC.SUCCESS
45
46
47class DictHandler(socketserver.BaseRequestHandler):
48    """Handler class for DICT connections.
49
50    """
51    def handle(self):
52        """
53        Simple function which responds to all queries with a 552.
54        """
55        try:
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"))
59
60            # Receive the request.
61            data = self.request.recv(1024).strip()
62            log.debug("[DICT] Incoming data: %r", data)
63
64            if VERIFIED_REQ in data:
65                log.debug("[DICT] Received verification request from test "
66                          "framework")
67                response_data = VERIFIED_RSP.format(pid=os.getpid())
68            else:
69                log.debug("[DICT] Received normal request")
70                response_data = "No matches"
71
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"))
76
77        except IOError:
78            log.exception("[DICT] IOError hit during request")
79
80
81def get_options():
82    parser = argparse.ArgumentParser()
83
84    parser.add_argument("--port", action="store", default=9016,
85                        type=int, help="port to listen on")
86    parser.add_argument("--host", action="store", default=HOST,
87                        help="host to listen on")
88    parser.add_argument("--verbose", action="store", type=int, default=0,
89                        help="verbose output")
90    parser.add_argument("--pidfile", action="store",
91                        help="file name for the PID")
92    parser.add_argument("--logfile", action="store",
93                        help="file name for the log")
94    parser.add_argument("--srcdir", action="store", help="test directory")
95    parser.add_argument("--id", action="store", help="server ID")
96    parser.add_argument("--ipv4", action="store_true", default=0,
97                        help="IPv4 flag")
98
99    return parser.parse_args()
100
101
102def setup_logging(options):
103    """
104    Set up logging from the command line options
105    """
106    root_logger = logging.getLogger()
107    add_stdout = False
108
109    formatter = logging.Formatter("%(asctime)s %(levelname)-5.5s %(message)s")
110
111    # Write out to a logfile
112    if options.logfile:
113        handler = logging.FileHandler(options.logfile, mode="w")
114        handler.setFormatter(formatter)
115        handler.setLevel(logging.DEBUG)
116        root_logger.addHandler(handler)
117    else:
118        # The logfile wasn't specified. Add a stdout logger.
119        add_stdout = True
120
121    if options.verbose:
122        # Add a stdout logger as well in verbose mode
123        root_logger.setLevel(logging.DEBUG)
124        add_stdout = True
125    else:
126        root_logger.setLevel(logging.INFO)
127
128    if add_stdout:
129        stdout_handler = logging.StreamHandler(sys.stdout)
130        stdout_handler.setFormatter(formatter)
131        stdout_handler.setLevel(logging.DEBUG)
132        root_logger.addHandler(stdout_handler)
133
134
135class ScriptRC(object):
136    """Enum for script return codes"""
137    SUCCESS = 0
138    FAILURE = 1
139    EXCEPTION = 2
140
141
142class ScriptException(Exception):
143    pass
144
145
146if __name__ == '__main__':
147    # Get the options from the user.
148    options = get_options()
149
150    # Setup logging using the user options
151    setup_logging(options)
152
153    # Run main script.
154    try:
155        rc = dictserver(options)
156    except Exception as e:
157        log.exception(e)
158        rc = ScriptRC.EXCEPTION
159
160    log.info("[DICT] Returning %d", rc)
161    sys.exit(rc)
162