Left: | ||
Right: |
OLD | NEW |
---|---|
(Empty) | |
1 #!/usr/bin/env python3 | |
2 | |
3 import argparse | |
4 import re | |
5 import sys | |
6 import threading | |
7 | |
8 from http.server import BaseHTTPRequestHandler, HTTPServer | |
9 from string import Template | |
10 | |
11 PROXY_TOKEN = 'http_' | |
12 REGEX = '\\$' + PROXY_TOKEN + '[a-z_0-9]+\\b' | |
13 DEFAULT_LOG = '$remote_addr - - [$time_local] "$request" $status $bytes_sent' | |
14 _lock = threading.Lock() | |
15 | |
16 | |
17 class Handler(BaseHTTPRequestHandler): | |
18 def get_header_values(self): | |
19 values = {} | |
20 headers = re.findall(REGEX, self.format) | |
21 for name in headers: | |
22 new_var = name[len(PROXY_TOKEN)+1:] | |
23 values[name[1:]] = self.headers.get(new_var, '-') | |
24 return values | |
25 | |
26 def write_info(self, args): | |
27 _lock.acquire() | |
28 message = Template(self.format) | |
Vasily Kuznetsov
2017/08/11 17:31:47
You should probably use try/finally here with the
| |
29 self.info_file.write(message.safe_substitute(args)) | |
30 if self.info_file is sys.stdout: | |
31 self.info_file.flush() | |
32 _lock.release() | |
33 | |
34 def do_POST(self): | |
35 status = 200 | |
36 content = bytes(self.response, 'UTF-8') | |
37 values = { | |
38 'remote_addr': self.address_string(), | |
39 'time_local': self.log_date_time_string(), | |
40 'request': self.requestline, | |
41 'status': status, | |
42 'bytes_sent': len(content), | |
43 } | |
44 values.update(self.get_header_values()) | |
45 self.write_info(values) | |
46 self.send_response(status) | |
47 self.send_header('Content-Type', values['content_type']) | |
48 self.send_header('Content-Length', values['content_length']) | |
49 self.end_headers() | |
50 self.wfile.write(content) | |
51 | |
52 | |
53 if __name__ == '__main__': | |
54 parser = argparse.ArgumentParser() | |
55 parser.add_argument('--port', action='store', | |
56 default=8000, type=int, | |
57 nargs='?', | |
58 help='Port to use [default: 8000]') | |
59 parser.add_argument('--response', action='store', | |
60 type=str, nargs='?', default='OK', | |
61 help='The response send to the client') | |
62 parser.add_argument('--format', action='store', | |
63 type=str, nargs='?', | |
64 default=DEFAULT_LOG, | |
65 help='Format of the log ouput') | |
66 parser.add_argument('info_file', action='store', | |
Vasily Kuznetsov
2017/08/11 17:31:47
I'm not sure if "info_file" is the best name, it s
| |
67 type=str, nargs='?', default='-', | |
68 help='The file where the logs will be written') | |
69 args = parser.parse_args() | |
70 setattr(Handler, 'format', args.format) | |
71 setattr(Handler, 'response', args.response) | |
72 server_address = ('', args.port) | |
73 httpd = HTTPServer(server_address, Handler) | |
74 if args.info_file and args.info_file != '-': | |
75 fh = open(args.info_file, 'a') | |
76 else: | |
77 fh = sys.stdout | |
78 setattr(Handler, 'info_file', fh) | |
Vasily Kuznetsov
2017/08/11 17:31:47
I think it would be more clear to order the steps
| |
79 try: | |
80 httpd.serve_forever() | |
81 except: | |
82 if args.info_file != '-': | |
83 fh.close() | |
OLD | NEW |