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