#!/usr/bin/env python3
"""HTTP server for the MidiManagerWinrt UAF POC.

Serves poc.html from this directory and MojoJS bindings from the Chromium
out/<build>/gen tree under /gen/.

Usage:
  python3 server.py [--port 8080] [--gen <path-to-out/build/gen>]
"""
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/fragment.
      path = path.split('?', 1)[0].split('#', 1)[0]
      if path.startswith('/gen/'):
        rel = os.path.normpath(path[len('/gen/'):])
        # Prevent path traversal out of gen_dir.
        if rel.startswith('..'):
          return '/dev/null'
        return os.path.join(gen_dir, rel)
      # Default: serve from the POC directory.
      rel = os.path.normpath(path.lstrip('/'))
      if rel.startswith('..'):
        return '/dev/null'
      return os.path.join(POC_DIR, rel or 'poc.html')

    def log_message(self, fmt, *args):
      sys.stderr.write('[server] %s\n' % (fmt % args))

  return Handler


def main():
  ap = argparse.ArgumentParser()
  ap.add_argument('--port', type=int, default=8080)
  ap.add_argument(
      '--gen',
      default=os.environ.get('CHROMIUM_GEN_DIR',
                             '/path/to/chromium/src/out/asan/gen'),
      help='Path to out/<build>/gen for MojoJS bindings')
  args = ap.parse_args()

  if not os.path.isdir(args.gen):
    sys.stderr.write('[!] gen dir not found: %s\n' % args.gen)
    sys.stderr.write('    pass --gen <chromium>/src/out/<build>/gen\n')

  socketserver.TCPServer.allow_reuse_address = True
  with socketserver.TCPServer(('127.0.0.1', args.port),
                              make_handler(args.gen)) as httpd:
    sys.stderr.write('[+] serving on http://127.0.0.1:%d/poc.html\n' %
                     args.port)
    sys.stderr.write('[+] gen dir: %s\n' % args.gen)
    httpd.serve_forever()


if __name__ == '__main__':
  main()