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