dictserver.py 6.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197
  1. #!/usr/bin/env python3
  2. # -*- coding: utf-8 -*-
  3. #***************************************************************************
  4. # _ _ ____ _
  5. # Project ___| | | | _ \| |
  6. # / __| | | | |_) | |
  7. # | (__| |_| | _ <| |___
  8. # \___|\___/|_| \_\_____|
  9. #
  10. # Copyright (C) 2008 - 2022, 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. from __future__ import (absolute_import, division, print_function,
  29. unicode_literals)
  30. import argparse
  31. import logging
  32. import os
  33. import sys
  34. from util import ClosingFileHandler
  35. try: # Python 2
  36. import SocketServer as socketserver
  37. except ImportError: # Python 3
  38. import socketserver
  39. log = logging.getLogger(__name__)
  40. HOST = "localhost"
  41. # The strings that indicate the test framework is checking our aliveness
  42. VERIFIED_REQ = b"verifiedserver"
  43. VERIFIED_RSP = "WE ROOLZ: {pid}"
  44. def dictserver(options):
  45. """
  46. Starts up a TCP server with a DICT handler and serves DICT requests
  47. forever.
  48. """
  49. if options.pidfile:
  50. pid = os.getpid()
  51. # see tests/server/util.c function write_pidfile
  52. if os.name == "nt":
  53. pid += 65536
  54. with open(options.pidfile, "w") as f:
  55. f.write(str(pid))
  56. local_bind = (options.host, options.port)
  57. log.info("[DICT] Listening on %s", local_bind)
  58. # Need to set the allow_reuse on the class, not on the instance.
  59. socketserver.TCPServer.allow_reuse_address = True
  60. server = socketserver.TCPServer(local_bind, DictHandler)
  61. server.serve_forever()
  62. return ScriptRC.SUCCESS
  63. class DictHandler(socketserver.BaseRequestHandler):
  64. """Handler class for DICT connections.
  65. """
  66. def handle(self):
  67. """
  68. Simple function which responds to all queries with a 552.
  69. """
  70. try:
  71. # First, send a response to allow the server to continue.
  72. rsp = "220 dictserver <xnooptions> <msgid@msgid>\n"
  73. self.request.sendall(rsp.encode("utf-8"))
  74. # Receive the request.
  75. data = self.request.recv(1024).strip()
  76. log.debug("[DICT] Incoming data: %r", data)
  77. if VERIFIED_REQ in data:
  78. log.debug("[DICT] Received verification request from test "
  79. "framework")
  80. pid = os.getpid()
  81. # see tests/server/util.c function write_pidfile
  82. if os.name == "nt":
  83. pid += 65536
  84. response_data = VERIFIED_RSP.format(pid=pid)
  85. else:
  86. log.debug("[DICT] Received normal request")
  87. response_data = "No matches"
  88. # Send back a failure to find.
  89. response = "552 {0}\n".format(response_data)
  90. log.debug("[DICT] Responding with %r", response)
  91. self.request.sendall(response.encode("utf-8"))
  92. except IOError:
  93. log.exception("[DICT] IOError hit during request")
  94. def get_options():
  95. parser = argparse.ArgumentParser()
  96. parser.add_argument("--port", action="store", default=9016,
  97. type=int, help="port to listen on")
  98. parser.add_argument("--host", action="store", default=HOST,
  99. help="host to listen on")
  100. parser.add_argument("--verbose", action="store", type=int, default=0,
  101. help="verbose output")
  102. parser.add_argument("--pidfile", action="store",
  103. help="file name for the PID")
  104. parser.add_argument("--logfile", action="store",
  105. help="file name for the log")
  106. parser.add_argument("--srcdir", action="store", help="test directory")
  107. parser.add_argument("--id", action="store", help="server ID")
  108. parser.add_argument("--ipv4", action="store_true", default=0,
  109. help="IPv4 flag")
  110. return parser.parse_args()
  111. def setup_logging(options):
  112. """
  113. Set up logging from the command line options
  114. """
  115. root_logger = logging.getLogger()
  116. add_stdout = False
  117. formatter = logging.Formatter("%(asctime)s %(levelname)-5.5s %(message)s")
  118. # Write out to a logfile
  119. if options.logfile:
  120. handler = ClosingFileHandler(options.logfile)
  121. handler.setFormatter(formatter)
  122. handler.setLevel(logging.DEBUG)
  123. root_logger.addHandler(handler)
  124. else:
  125. # The logfile wasn't specified. Add a stdout logger.
  126. add_stdout = True
  127. if options.verbose:
  128. # Add a stdout logger as well in verbose mode
  129. root_logger.setLevel(logging.DEBUG)
  130. add_stdout = True
  131. else:
  132. root_logger.setLevel(logging.INFO)
  133. if add_stdout:
  134. stdout_handler = logging.StreamHandler(sys.stdout)
  135. stdout_handler.setFormatter(formatter)
  136. stdout_handler.setLevel(logging.DEBUG)
  137. root_logger.addHandler(stdout_handler)
  138. class ScriptRC(object):
  139. """Enum for script return codes"""
  140. SUCCESS = 0
  141. FAILURE = 1
  142. EXCEPTION = 2
  143. class ScriptException(Exception):
  144. pass
  145. if __name__ == '__main__':
  146. # Get the options from the user.
  147. options = get_options()
  148. # Setup logging using the user options
  149. setup_logging(options)
  150. # Run main script.
  151. try:
  152. rc = dictserver(options)
  153. except Exception as e:
  154. log.exception(e)
  155. rc = ScriptRC.EXCEPTION
  156. if options.pidfile and os.path.isfile(options.pidfile):
  157. os.unlink(options.pidfile)
  158. log.info("[DICT] Returning %d", rc)
  159. sys.exit(rc)