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