HTTP echo server supporting msgpack

agent.py

#!/usr/bin/env python

# Simple echo-like HTTP server that supports echoing
# of msgpack payloads.
#
# Ideally used to debug datadog tracing messages
#
# Dependencies:
#   - msgpack (`pip install msgpack`)

from http.server import BaseHTTPRequestHandler, HTTPServer

import argparse
import logging

import msgpack


class _Server(BaseHTTPRequestHandler):
    def __init__(self, show_headers=False):
        self.show_headers = show_headers

    def __call__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

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

    def do_GET(self):
        if self.show_headers:
            logging.info('GET request,\nPath: %s\nHeaders:\n%s', str(self.path), str(self.headers))
        else:
            logging.info('GET request,\nPath: %s', str(self.path), str(self.headers))
        self._set_response()
        self.wfile.write('GET request for {}'.format(self.path).encode('utf-8'))

    def do_POST(self):
        content_length = int(self.headers['Content-Length'])
        content_type = self.headers['Content-Type']
        post_data = self.rfile.read(content_length)
        if self.show_headers:
            logging.info('POST request,\nPath: %s\nHeaders:\n%s\n\nBody:\n%s',
                    str(self.path), str(self.headers), _parse(post_data, content_type))
        else:
            logging.info('POST request,\nPath: %s\nBody:\n%s',
                    str(self.path), _parse(post_data, content_type))


        self._set_response()
        self.wfile.write('POST request for {}'.format(self.path).encode('utf-8'))

    def do_PUT(self):
        content_length = int(self.headers['Content-Length'])
        content_type = self.headers['Content-Type']
        put_data = self.rfile.read(content_length)
        if self.show_headers:
            logging.info('PUT request,\nPath: %s\nHeaders:\n%s\n\nBody:\n%s',
                    str(self.path), str(self.headers), _parse(put_data, content_type))
        else:
            logging.info('POST request,\nPath: %s\nBody:\n%s',
                    str(self.path), _parse(put_data, content_type))

        self._set_response()
        self.wfile.write('PUT request for {}'.format(self.path).encode('utf-8'))


def _parse(content, content_type):
    if 'msgpack' in content_type:
        parsed = msgpack.unpackb(content)
        return str(parsed)
    return content.decode('utf-8')


def run(handler, port=8126):
    logging.basicConfig(level=logging.INFO)
    server_address = ('', port)
    httpd = HTTPServer(server_address, handler)
    logging.info('start server at %d...', port)
    try:
        httpd.serve_forever()
    except KeyboardInterrupt:
        pass
    httpd.server_close()
    logging.info('stop server...')


if __name__ == '__main__':
    PARSER = argparse.ArgumentParser(description='HTTP echo server supporting msgpack')
    PARSER.add_argument('-p', '--port', type=int, default=8126, help='port to listen on')
    PARSER.add_argument('-H', '--headers', action='store_true', help='output request headers')

    ARGS = PARSER.parse_args()
    run(_Server(ARGS.headers), port=ARGS.port)

agent.py