tests/dictserver.py

   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)