#!/usr/bin/env python3
"""HTTP server for the WebUSB SetInterfaceAlternateSetting bypass POC.

Serves poc.html from this directory and the generated MojoJS lite bindings
from a Chromium build directory under /gen/.

Usage:
    python3 server.py [--port 8080] [--build-dir /path/to/chromium/src/out/asan]
"""

import argparse
import http.server
import os
import socketserver
import sys

POC_DIR = os.path.dirname(os.path.abspath(__file__))


def make_handler(gen_dir):
    class Handler(http.server.SimpleHTTPRequestHandler):
        def translate_path(self, path):
            # Strip query string / fragment.
            p = path.split('?', 1)[0].split('#', 1)[0]
            if p.startswith('/gen/'):
                rel = os.path.normpath(p[len('/gen/'):])
                if rel.startswith('..'):
                    return '/dev/null'
                return os.path.join(gen_dir, rel)
            # Everything else from the POC directory.
            rel = os.path.normpath(p.lstrip('/'))
            if rel.startswith('..'):
                return '/dev/null'
            if rel in ('', '.'):
                rel = 'poc.html'
            return os.path.join(POC_DIR, rel)

        def end_headers(self):
            self.send_header('Cache-Control', 'no-store')
            super().end_headers()

    return Handler


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--port', type=int, default=8080)
    ap.add_argument('--build-dir',
                    default='/chromium/chromium/src/out/asan',
                    help='Chromium build dir containing gen/ with mojom-lite '
                         'bindings')
    args = ap.parse_args()

    gen_dir = os.path.join(args.build_dir, 'gen')
    if not os.path.isdir(gen_dir):
        sys.stderr.write(f'warning: {gen_dir} does not exist; /gen/ paths '
                         f'will 404\n')

    socketserver.TCPServer.allow_reuse_address = True
    with socketserver.TCPServer(('127.0.0.1', args.port),
                                make_handler(gen_dir)) as httpd:
        sys.stderr.write(
            f'serving POC on http://localhost:{args.port}/poc.html\n')
        httpd.serve_forever()


if __name__ == '__main__':
    main()