#!/usr/bin/env python3
"""Local HTTP adapter for Courier package downloads and Git HTTPS operations."""

import argparse
import gzip
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
import io
import re
import shlex
import subprocess
import zlib


MAX_GIT_BODY_SIZE = 1024 * 1024
MAX_GIT_PUSH_BODY_SIZE = 16 * 1024 * 1024


def build_command(method, target, courier_cli='courier-cli', *, git_protocol=None):
    """Translate /HOST/CREDENTIAL_ID/PATH?QUERY into a subprocess argument list."""
    if (
        not target.startswith('/') or '#' in target
        or any(ord(c) <= 32 or ord(c) == 127 for c in target)
    ):
        raise ValueError('invalid request target')
    git = target.startswith('/_git/')
    if git:
        target = target[len('/_git'):]
    elif method not in ('GET', 'HEAD') or git_protocol is not None:
        raise ValueError('registry routes support only GET and HEAD')
    route, separator, query = target.partition('?')
    parts = route[1:].split('/', 2)
    if len(parts) < 2:
        raise ValueError('expected /HOST/CREDENTIAL_ID/PATH')
    host, credential = parts[:2]
    if not re.fullmatch(r'[A-Za-z0-9.-]+(?::[0-9]+)?', host):
        raise ValueError('invalid remote host')
    if not re.fullmatch(r'[A-Z][A-Z0-9_]{0,63}', credential):
        raise ValueError('invalid credential ID')
    path = '/' + (parts[2] if len(parts) == 3 else '')
    if git:
        segments = path[1:].split('/')
        if any(c in path for c in '%\\') or any(part in ('', '.', '..') for part in segments):
            raise ValueError('Git repository path must be canonical')
        discovery = (method == 'GET' and path.endswith('.git/info/refs')
                     and query in ('service=git-upload-pack', 'service=git-receive-pack'))
        operation = (method == 'POST' and not separator
                     and path.endswith(('.git/git-upload-pack', '.git/git-receive-pack')))
        if not (discovery or operation):
            raise ValueError('only Git discovery, fetch and push endpoints are supported')
        if git_protocol not in (None, 'version=2'):
            raise ValueError('only the version=2 Git protocol header is supported')
    command = [
        courier_cli, 'request', '--stream-http', '-X', method,
        '-H', f'{"Authorization" if git else "PRIVATE-TOKEN"}: ${{{credential}}}',
        f'https://{host}{path}{separator}{query}',
    ]
    if git_protocol is not None:
        command.extend(['--git-protocol', git_protocol])
    if method == 'POST':
        command.extend(['--data-file', '-'])
    return command


class ProxyHandler(BaseHTTPRequestHandler):
    def setup(self):
        super().setup()
        self.connection.settimeout(30)

    def parse_request(self):
        parsed = super().parse_request()
        if parsed and self.server.verbose:
            self.log_message('received %s %s', self.command, self.path)
        return parsed

    def log_request(self, code='-', size='-'):
        # Responses are usually relayed directly; log arrivals consistently instead.
        pass

    def do_GET(self):
        self.close_connection = True
        try:
            protocols = self.headers.get_all('Git-Protocol', []) if self.path.startswith('/_git/') else []
            if len(protocols) > 1:
                raise ValueError('duplicate Git-Protocol header')
            command = build_command(self.command, self.path, self.server.courier_cli,
                                    git_protocol=protocols[0] if protocols else None)
            body = self.read_body()
        except (ValueError, OSError, EOFError, zlib.error) as error:
            self.send_error(400, str(error))
            return
        if self.server.dry_run:
            print(shlex.join(command), flush=True)
            body = b'200 OK\n'
            self.send_response(200)
            self.send_header('Content-Type', 'text/plain')
            self.send_header('Content-Length', str(len(body)))
            self.end_headers()
            if self.command != 'HEAD':
                self.wfile.write(body)
            return
        try:
            process = subprocess.Popen(command, stdin=subprocess.PIPE if body else subprocess.DEVNULL,
                                       stdout=subprocess.PIPE)
        except OSError:
            self.send_error(502, 'could not start courier-cli')
            return

        sent = False
        try:
            if body:
                process.stdin.write(body)
                process.stdin.close()
            while block := process.stdout.read1(65536):
                self.wfile.write(block)
                sent = True
            status = process.wait()
            if not sent:
                self.send_error(502, 'courier-cli returned no HTTP response')
            elif status:
                self.log_error('courier-cli exited with status %s', status)
        except (BrokenPipeError, ConnectionResetError):
            if not sent:
                self.send_error(502, 'courier-cli did not accept the request')
        finally:
            if process.stdin is not None:
                try:
                    process.stdin.close()
                except BrokenPipeError:
                    pass
            process.stdout.close()
            if process.poll() is None:
                process.terminate()
                try:
                    process.wait(timeout=5)
                except subprocess.TimeoutExpired:
                    process.kill()
                    process.wait()

    do_HEAD = do_GET
    do_POST = do_GET

    def read_body(self):
        lengths = self.headers.get_all('Content-Length', [])
        transfer = self.headers.get_all('Transfer-Encoding', [])
        push = self.path.split('?', 1)[0].endswith('.git/git-receive-pack')
        chunked = push and self.command == 'POST' and [value.lower() for value in transfer] == ['chunked']
        if len(lengths) > 1 or transfer and (lengths or not chunked):
            raise ValueError('ambiguous or unsupported request framing')
        if lengths and not re.fullmatch(r'[0-9]+', lengths[0]):
            raise ValueError('invalid content length')
        length = int(lengths[0]) if lengths else 0
        if self.command != 'POST':
            if length != 0:
                raise ValueError('request bodies are not supported')
            return None
        maximum = MAX_GIT_PUSH_BODY_SIZE if push else MAX_GIT_BODY_SIZE
        service = 'git-receive-pack' if push else 'git-upload-pack'
        if self.headers.get_all('Content-Type', []) != [f'application/x-{service}-request']:
            raise ValueError('unsupported Git request content type')
        if chunked:
            body = self.read_chunks(maximum)
        else:
            if not 0 < length <= maximum:
                raise ValueError(f'Git body must contain 1 to {maximum} bytes')
            body = self.rfile.read(length)
            if len(body) != length:
                raise ValueError('incomplete Git request body')
        encoding = self.headers.get_all('Content-Encoding', [])
        if encoding == ['gzip']:
            with gzip.GzipFile(fileobj=io.BytesIO(body)) as compressed:
                body = compressed.read(maximum + 1)
        elif encoding not in ([], ['identity']):
            raise ValueError('unsupported Git request content encoding')
        if not 0 < len(body) <= maximum:
            raise ValueError('decoded Git body is too large or empty')
        return body

    def read_chunks(self, maximum):
        """Accept Git's chunked pushes, without extensions, trailers or excess data."""
        body = bytearray()
        while True:
            line = self.rfile.readline(32)
            if not re.fullmatch(rb'[0-9a-fA-F]{1,8}\r\n', line):
                raise ValueError('invalid Git chunk size')
            size = int(line, 16)
            if size == 0:
                if self.rfile.read(2) != b'\r\n':
                    raise ValueError('Git chunk trailers are not supported')
                return bytes(body)
            if size > maximum - len(body):
                raise ValueError('Git chunked body is too large')
            block = self.rfile.read(size)
            if len(block) != size or self.rfile.read(2) != b'\r\n':
                raise ValueError('incomplete Git chunk')
            body.extend(block)


class ProxyServer(ThreadingHTTPServer):
    def __init__(self, address, courier_cli='courier-cli', *, dry_run=False, verbose=False):
        self.courier_cli = courier_cli
        self.dry_run = dry_run
        self.verbose = verbose
        super().__init__(address, ProxyHandler)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--port', type=int, default=18080)
    parser.add_argument('--courier-cli', default='courier-cli')
    parser.add_argument('-v', '--verbose', action='store_true',
                        help='log one line to stderr when each request arrives')
    parser.add_argument('--dry-run', action='store_true',
                        help='print commands and return 200 OK without executing courier-cli')
    args = parser.parse_args()
    with ProxyServer(('127.0.0.1', args.port), args.courier_cli,
                     dry_run=args.dry_run, verbose=args.verbose) as server:
        try:
            server.serve_forever()
        except KeyboardInterrupt:
            pass


if __name__ == '__main__':
    main()
