echod.py

#!/usr/bin/env python3
from argparse import ArgumentParser
import logging
from pathlib import Path
import http.server
import socketserver
from uuid import uuid4
from functools import wraps

THIS_DIR = Path(__file__).parent.absolute()

def auto_log(method):
    @wraps(method)
    def wrapper(self, *args, **kwargs):
        request_log = [
            '# %s' % self.request_id,
            ' '.join([self.command, self.path, self.protocol_version]),
            str(self.headers),
        ]
        if self.command in ['POST', 'PUT']:
            body = self.read_body()
            if isinstance(body, bytes):
                body = body.decode('utf8')
            request_log.append('########################')
            request_log.append(body)
        request_log.append('------------------------')
        logging.info('\n'.join(request_log))
        return method(self, *args, **kwargs)
    return wrapper

class EchoHandler(http.server.BaseHTTPRequestHandler):
    def __init__(self, *args, **kwargs):
        self.request_id = str(uuid4())
        self.request_body = None
        super(EchoHandler, self).__init__(*args, **kwargs)

    def send_ok(self):
        self.send_header('content-type', 'text/plain')
        self.send_header('content-length', '0')
        self.send_response(200)
        self.end_headers()

    @auto_log
    def do_GET(self):
        self.send_ok()

    @auto_log
    def do_POST(self):
        self.send_ok()

    def read_body(self):
        if self.command not in ['POST', 'PUT']:
            return None
        if self.request_body is not None:
            return self.request_body
        clen = int(self.headers.get('content-length', 0))
        max_len = 1024 * 1024 * 32 # 32 M
        assert(clen < max_len)
        self.request_body = self.rfile.read(clen)
        return self.request_body


def main():
    log_filename = THIS_DIR.joinpath('logs-echod.txt')
    log_format = '[%(asctime)s %(filename)s:%(lineno)d] %(message)s'
    logging.basicConfig(level=logging.DEBUG, format=log_format, filename=log_filename)
    # logging.basicConfig(level=logging.INFO, format=log_format)
    socketserver.TCPServer.allow_reuse_address = True

    cli_parser = ArgumentParser()
    cli_parser.add_argument('--port', default=13579, type=int)
    cli_parser.add_argument('--address', default='0.0.0.0', type=str)
    cli_configure = cli_parser.parse_args()
    server_address = (cli_configure.address, cli_configure.port)
    httpd = socketserver.TCPServer(server_address, EchoHandler)
    logging.info('server is running on %s', server_address)
    try:
        httpd.serve_forever()
    except KeyboardInterrupt:
        httpd.server_close()


if __name__ == '__main__':
    main()