#!/usr/bin/env python3
"""Ephemeral encrypted rendezvous mailbox. Python 3.11+, no dependencies.

The native clients generate the encryption key. This process receives only
opaque ciphertext and a SEPARATE mailbox capability. Nothing is persisted.
"""
import argparse
import hashlib
import hmac
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
import re
import socket
import ssl
import threading
import time

MAX_BODY = 1600 * 1024
MAX_MEMORY = 128 * 1024 * 1024
MAX_ROOMS = 32
TTL = 900
REQUEST_DEADLINE = 15
ROUTE = re.compile(r"/v1/rooms/([0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12})/(offer-i|offer-r|ack-i|ack-r)\Z")
TOKEN = re.compile(r"Bearer ([A-Za-z0-9_-]{43}=?)\Z")


class Mailboxes:
    def __init__(self, clock=time.monotonic):
        self.clock = clock
        self.lock = threading.Lock()
        self.rooms = {}

    def _expire(self, now):
        for key in list(self.rooms):
            if self.rooms[key][0] <= now:
                del self.rooms[key]

    def expire(self):
        with self.lock:
            self._expire(self.clock())

    def exchange(self, room, slot, token, body=None):
        with self.lock:
            now = self.clock()
            self._expire(now)
            record = self.rooms.get(room)
            token_hash = hashlib.sha256(token.encode("ascii")).digest()
            if record is None:
                if body is None:
                    return 404, b""
                if slot != "offer-i":
                    return 404, b""
                if len(self.rooms) >= MAX_ROOMS:
                    return 503, b""
                record = (now + TTL, token_hash, {})
                self.rooms[room] = record
            if not hmac.compare_digest(record[1], token_hash):
                return 403, b""
            slots = record[2]
            if body is None:
                return (200, slots[slot]) if slot in slots else (404, b"")
            if slot in slots:
                return (204, b"") if hmac.compare_digest(slots[slot], body) else (409, b"")
            total = sum(len(value) for r in self.rooms.values() for value in r[2].values())
            if total + len(body) > MAX_MEMORY:
                return 503, b""
            slots[slot] = body
            return 204, b""


class RelayServer(ThreadingHTTPServer):
    daemon_threads = True
    request_queue_size = 16

    def __init__(self, address, boxes=None):
        self.boxes = boxes or Mailboxes()
        self.capacity = threading.BoundedSemaphore(16)
        super().__init__(address, Handler)

    def get_request(self):
        client, address = super().get_request()
        client.settimeout(10)
        return client, address

    def process_request(self, request, client_address):
        if not self.capacity.acquire(blocking=False):
            self.shutdown_request(request)
            return
        try:
            super().process_request(request, client_address)
        except BaseException:
            self.capacity.release()
            raise

    def process_request_thread(self, request, client_address):
        def expire_request():
            try:
                request.shutdown(socket.SHUT_RDWR)
            except OSError:
                pass
            request.close()
        deadline = threading.Timer(REQUEST_DEADLINE, expire_request)
        deadline.daemon = True
        deadline.start()
        try:
            super().process_request_thread(request, client_address)
        finally:
            deadline.cancel()
            self.capacity.release()

    def service_actions(self):
        # Expire idle ciphertext even when no later request arrives.
        self.boxes.expire()

    def handle_error(self, request, client_address):
        # No paths, capabilities, network addresses or ciphertext in logs.
        pass


class Handler(BaseHTTPRequestHandler):
    protocol_version = "HTTP/1.0"
    server_version = "IroncladRendezvous/1"
    sys_version = ""

    def log_message(self, *_):
        pass

    def respond(self, status, body=b""):
        self.send_response(status)
        self.send_header("Content-Length", str(len(body)))
        self.send_header("Content-Type", "application/octet-stream")
        self.send_header("Cache-Control", "no-store")
        self.send_header("X-Content-Type-Options", "nosniff")
        self.send_header("Connection", "close")
        self.end_headers()
        if body:
            self.wfile.write(body)
        self.close_connection = True

    def handle_mailbox(self, write):
        route = ROUTE.fullmatch(self.path)
        auths = self.headers.get_all("Authorization", [])
        auth = TOKEN.fullmatch(auths[0]) if len(auths) == 1 else None
        if route is None or auth is None:
            self.respond(400)
            return
        lengths = self.headers.get_all("Content-Length", [])
        if self.headers.get_all("Transfer-Encoding") or (write and len(lengths) != 1):
            self.respond(400)
            return
        if write:
            if not re.fullmatch(r"[1-9][0-9]{0,7}", lengths[0]):
                self.respond(400)
                return
            size = int(lengths[0])
            if size < 29 or size > MAX_BODY:
                self.respond(413)
                return
            try:
                body = self.rfile.read(size)
            except (TimeoutError, socket.timeout):
                self.respond(408)
                return
            if len(body) != size:
                self.respond(400)
                return
        else:
            if lengths and lengths != ["0"]:
                self.respond(400)
                return
            body = None
        self.respond(*self.server.boxes.exchange(*route.groups(), auth[1], body))

    def do_GET(self):
        self.handle_mailbox(False)

    def do_PUT(self):
        self.handle_mailbox(True)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--bind", default="127.0.0.1")
    parser.add_argument("--port", type=int, default=8765)
    parser.add_argument("--cert", help="TLS certificate PEM")
    parser.add_argument("--key", help="TLS private key PEM")
    parser.add_argument("--loopback-http", action="store_true", help="Local QA or local TLS reverse proxy only")
    args = parser.parse_args()
    if args.loopback_http:
        if args.bind != "127.0.0.1" or args.cert or args.key:
            parser.error("Plain HTTP is allowed only on 127.0.0.1 without TLS options")
    elif not args.cert or not args.key:
        parser.error("Provide --cert and --key for HTTPS, or explicit loopback QA mode")
    server = RelayServer((args.bind, args.port))
    if not args.loopback_http:
        context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
        context.minimum_version = ssl.TLSVersion.TLSv1_2
        context.load_cert_chain(args.cert, args.key)
        # Handshake after accept in the bounded worker, not in the accept loop.
        original_get = server.get_request
        def tls_request():
            client, address = original_get()
            return context.wrap_socket(client, server_side=True, do_handshake_on_connect=False), address
        server.get_request = tls_request
    print("Encrypted rendezvous mailbox running; no request logging or disk storage.", flush=True)
    try:
        server.serve_forever(poll_interval=0.5)
    except KeyboardInterrupt:
        pass
    finally:
        server.server_close()


if __name__ == "__main__":
    main()
