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