-
Notifications
You must be signed in to change notification settings - Fork 39
/
run.py
129 lines (99 loc) · 3.56 KB
/
run.py
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
import fcntl
import argparse
import sys
import logging
from vaurien.server import DoWeirdThingsPlease
from vaurien.config import load_into_settings, DEFAULT_SETTINGS
from vaurien import __version__, logger
LOG_LEVELS = {
"critical": logging.CRITICAL,
"error": logging.ERROR,
"warning": logging.WARNING,
"info": logging.INFO,
"debug": logging.DEBUG}
LOG_FMT = r"%(asctime)s [%(process)d] [%(levelname)s] %(message)s"
LOG_DATE_FMT = r"%Y-%m-%d %H:%M:%S"
def close_on_exec(fd):
flags = fcntl.fcntl(fd, fcntl.F_GETFD)
flags |= fcntl.FD_CLOEXEC
fcntl.fcntl(fd, fcntl.F_SETFD, flags)
def configure_logger(logger, level='INFO', output="-"):
loglevel = LOG_LEVELS.get(level.lower(), logging.INFO)
logger.setLevel(loglevel)
if output == "-":
h = logging.StreamHandler()
else:
h = logging.FileHandler(output)
close_on_exec(h.stream.fileno())
fmt = logging.Formatter(LOG_FMT, LOG_DATE_FMT)
h.setFormatter(fmt)
logger.addHandler(h)
def get_statsd_from_settings(settings):
if settings['enabled']:
from statsd import StatsdClient
statsd = StatsdClient(host=settings['host'],
port=settings['port'],
prefix=settings['prefix'],
sample_rate=settings['sample_rate'])
else:
statsd = None
return statsd
def main():
parser = argparse.ArgumentParser(description='Runs a mean TCP proxy.')
# other arguments
parser.add_argument('--config', help='Configuration file', default=None)
parser.add_argument('--version', action='store_true', default=False,
help='Displays version and exits.')
# get the values from the default config
keys = DEFAULT_SETTINGS.keys()
keys.sort()
for key in keys:
if key.startswith('vaurien'):
key = key[len('vaurien.'):]
parser.add_argument('--%s' % key, default=None)
parser.add_argument('--log-level', dest='loglevel', default='info',
choices=LOG_LEVELS.keys() + [key.upper() for key in
LOG_LEVELS.keys()],
help="log level")
parser.add_argument('--log-output', dest='logoutput', default='-',
help="log output")
args = parser.parse_args()
if args.version:
print(__version__)
sys.exit(0)
# configure the logger
configure_logger(logger, args.loglevel, args.logoutput)
settings = DEFAULT_SETTINGS.copy()
if args.config is not None:
# read the config if provided
try:
load_into_settings(args.config, settings)
except ValueError, e:
print(e)
sys.exit(1)
# and finally overwrite with the commandline arguments
for key in settings.keys():
prefix = ''
if key.startswith('vaurien'):
key = key[len('vaurien.'):]
prefix = 'vaurien.'
try:
value = getattr(args, key)
except AttributeError:
value = None
if value is not None:
settings[prefix + key] = value
statsd = get_statsd_from_settings(settings.getsection('statsd'))
# creating the server
server = DoWeirdThingsPlease(local=settings['vaurien.local'],
distant=settings['vaurien.distant'],
settings=settings, statsd=statsd,
logger=logger)
try:
server.serve_forever()
except KeyboardInterrupt:
sys.exit(0)
finally:
logger.info('Bye!')
if __name__ == '__main__':
main()