Home | History | Annotate | Download | only in tests 
      1 #!/usr/bin/env python
      2 # -*- coding: utf-8 -*-
      3 #
      4 """ DICT server """
      5 
      6 from __future__ import (absolute_import, division, print_function,
      7                         unicode_literals)
      8 import argparse
      9 import os
     10 import sys
     11 import logging
     12 try:  # Python 2
     13     import SocketServer as socketserver
     14 except ImportError:  # Python 3
     15     import socketserver
     16 
     17 
     18 log = logging.getLogger(__name__)
     19 HOST = "localhost"
     20 
     21 # The strings that indicate the test framework is checking our aliveness
     22 VERIFIED_REQ = b"verifiedserver"
     23 VERIFIED_RSP = "WE ROOLZ: {pid}"
     24 
     25 
     26 def 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 = (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 
     47 class 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 
     81 def 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("--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,
     95                         help="IPv4 flag")
     96 
     97     return parser.parse_args()
     98 
     99 
    100 def setup_logging(options):
    101     """
    102     Set up logging from the command line options
    103     """
    104     root_logger = logging.getLogger()
    105     add_stdout = False
    106 
    107     formatter = logging.Formatter("%(asctime)s %(levelname)-5.5s %(message)s")
    108 
    109     # Write out to a logfile
    110     if options.logfile:
    111         handler = logging.FileHandler(options.logfile, mode="w")
    112         handler.setFormatter(formatter)
    113         handler.setLevel(logging.DEBUG)
    114         root_logger.addHandler(handler)
    115     else:
    116         # The logfile wasn't specified. Add a stdout logger.
    117         add_stdout = True
    118 
    119     if options.verbose:
    120         # Add a stdout logger as well in verbose mode
    121         root_logger.setLevel(logging.DEBUG)
    122         add_stdout = True
    123     else:
    124         root_logger.setLevel(logging.INFO)
    125 
    126     if add_stdout:
    127         stdout_handler = logging.StreamHandler(sys.stdout)
    128         stdout_handler.setFormatter(formatter)
    129         stdout_handler.setLevel(logging.DEBUG)
    130         root_logger.addHandler(stdout_handler)
    131 
    132 
    133 class ScriptRC(object):
    134     """Enum for script return codes"""
    135     SUCCESS = 0
    136     FAILURE = 1
    137     EXCEPTION = 2
    138 
    139 
    140 class ScriptException(Exception):
    141     pass
    142 
    143 
    144 if __name__ == '__main__':
    145     # Get the options from the user.
    146     options = get_options()
    147 
    148     # Setup logging using the user options
    149     setup_logging(options)
    150 
    151     # Run main script.
    152     try:
    153         rc = dictserver(options)
    154     except Exception as e:
    155         log.exception(e)
    156         rc = ScriptRC.EXCEPTION
    157 
    158     log.info("[DICT] Returning %d", rc)
    159     sys.exit(rc)
    160