#!/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")


async def relay(request: web.Request) -> web.WebSocketResponse:
    """A WebSocket whose binary frames are piped to host:port untouched."""
    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()

    await asyncio.gather(pump_up(), pump_down())
    log.info("relay to %s:%d closed", host, port)
    return ws


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")
    args = parser.parse_args()

    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
    app = web.Application()
    app["web_dir"], app["relay_ports"] = args.dir, set(args.allow_ports)
    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())