#!/usr/bin/env -S uv run --quiet --script
# /// script
# requires-python = ">=3.11"
# dependencies = ["aiohttp>=3.9"]
# ///
"""gsmol bridge: serve the web client and relay WebSocket bytes to smolmaild TCP.

Browsers cannot open raw TCP sockets, so the Noise_NX session runs in the
browser and this bridge carries its bytes. The bridge holds no keys and never
parses a byte of the traffic; it is a pipe with a port allowlist, so it is not
an open proxy.

    uv run bridge.py                       # serve web/ and relay port 1961
    uv run bridge.py --allow-port 11961    # relay to an extra port
"""

from __future__ import annotations

import argparse
import asyncio
import logging
from pathlib import Path

import aiohttp
from aiohttp import web

log = logging.getLogger("gsmol")


def origin_allowed(request: web.Request) -> bool:
    """Whether this upgrade may open a relay.

    WebSockets are exempt from CORS, so without this check any page the user
    happens to visit could drive the relay to an allowlisted port on any host.
    Browsers always send Origin on an upgrade; non-browser clients (test/client.mjs
    on Node) send none, so an absent header is allowed and a foreign one is not.
    """
    origin = request.headers.get("Origin")
    if origin is None:
        return True
    own = {f"http://{request.host}", f"https://{request.host}"}
    return origin in own or origin in request.app["allow_origins"]


async def relay(request: web.Request) -> web.WebSocketResponse:
    """A WebSocket whose binary frames are piped to host:port untouched."""
    if not origin_allowed(request):
        raise web.HTTPForbidden(text="cross-origin WebSocket upgrade refused")
    port = int(request.match_info["port"])
    if port not in request.app["relay_ports"]:
        raise web.HTTPForbidden(text=f"port {port} is not allowed")
    host = request.match_info["host"]

    ws = web.WebSocketResponse(autoping=True)
    await ws.prepare(request)
    try:
        reader, writer = await asyncio.open_connection(host, port)
    except OSError as exc:
        await ws.close(code=1011, message=str(exc).encode())
        return ws
    log.info("relaying to %s:%d", host, port)

    async def pump_up() -> None:
        try:
            async for msg in ws:
                if msg.type == aiohttp.WSMsgType.BINARY:
                    writer.write(msg.data)
                    await writer.drain()
                elif msg.type == aiohttp.WSMsgType.ERROR:
                    break
        finally:
            writer.close()

    async def pump_down() -> None:
        try:
            while chunk := await reader.read(65536):
                await ws.send_bytes(chunk)
        except (ConnectionError, asyncio.IncompleteReadError):
            pass
        finally:
            await ws.close()

    # return_exceptions so a failure in one pump cannot leave the other running:
    # closing the writer ends the downstream read, closing the socket ends the
    # upstream iteration, so each side unblocks the other.
    await asyncio.gather(pump_up(), pump_down(), return_exceptions=True)
    try:
        await writer.wait_closed()
    except OSError:
        pass
    log.info("relay to %s:%d closed", host, port)
    return ws


# The seed lives in localStorage, so an XSS in this page is game over (README).
# The client loads no inline script, no inline style and no foreign origin, so a
# strict policy costs nothing. form-action stays 'self' rather than 'none' so the
# compose dialog's method="dialog" form is untouched.
#
# Nix normalizes every built file's mtime to the same fixed epoch for
# reproducibility, so Last-Modified is identical across every deploy — with no
# explicit Cache-Control, a browser's heuristic freshness calculation (based on
# that ~56-year-old timestamp) can treat a page as fresh indefinitely and never
# revalidate again, silently pinning a visitor to whatever JS they first
# loaded across every future deploy. no-cache forces a conditional GET each
# time rather than disabling caching outright — the ETag aiohttp already sends
# still turns an unchanged file into a bodyless 304.
@web.middleware
async def headers(request: web.Request, handler):
    response = await handler(request)
    response.headers.setdefault("Content-Security-Policy", "; ".join([
        "default-src 'self'",
        f"connect-src 'self' ws://{request.host} wss://{request.host}",
        "base-uri 'none'",
        "form-action 'self'",
        "object-src 'none'",
        "frame-ancestors 'none'",
    ]))
    response.headers.setdefault("Cache-Control", "no-cache")
    return response


async def index(request: web.Request) -> web.FileResponse:
    return web.FileResponse(Path(request.app["web_dir"]) / "index.html")


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--host", default="127.0.0.1")
    parser.add_argument("--port", type=int, default=8096)
    parser.add_argument("--dir", default=str(Path(__file__).parent / "web"))
    parser.add_argument("--allow-port", type=int, action="append", default=[1961],
                        metavar="PORT", dest="allow_ports",
                        help="relayed TCP port; repeatable")
    parser.add_argument("--allow-origin", action="append", default=[],
                        metavar="ORIGIN", dest="allow_origins",
                        help="extra browser origin permitted to open a relay; repeatable")
    args = parser.parse_args()

    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
    app = web.Application(middlewares=[headers])
    app["web_dir"], app["relay_ports"] = args.dir, set(args.allow_ports)
    app["allow_origins"] = set(args.allow_origins)
    app.router.add_get("/", index)
    app.router.add_get("/tcp/{host}/{port}", relay)
    app.router.add_static("/", args.dir, show_index=False)
    log.info("serving %s on http://%s:%d — allowed TCP ports: %s",
             args.dir, args.host, args.port, ", ".join(map(str, sorted(app["relay_ports"]))))
    web.run_app(app, host=args.host, port=args.port, print=None)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())