server.py

#!/usr/bin/env python3
"""AntScopeZ Web -- browser/iPhone bridge to AntScopeZ's Remote API.

Runs as its own process on the same machine as AntScopeZ. Connects
upstream to AntScopeZ's Remote API over 127.0.0.1 (same trust model
the Remote API itself uses) and serves a small web UI on 0.0.0.0:<port>
so a phone on the same Wi-Fi can reach it. See README.md.
"""

import argparse
import asyncio
import json
import logging
from pathlib import Path

from aiohttp import web, WSMsgType

from antscopez_client import AntScopeZBridge

logger = logging.getLogger("antscopez_web")

STATIC_DIR = Path(__file__).parent / "static"

# Commands a browser is allowed to trigger directly. subscribe/unsubscribe
# are deliberately excluded -- the bridge subscribes once for everyone
# up front (see AntScopeZBridge._connect_once) so a browser can't
# accidentally unsubscribe the shared stream for every other tab/phone.
ALLOWED_COMMANDS = {"status", "devices", "connect", "disconnect", "sweep", "stop", "last"}


async def ws_handler(request: web.Request) -> web.WebSocketResponse:
    ws = web.WebSocketResponse(heartbeat=30)
    await ws.prepare(request)

    bridge: AntScopeZBridge = request.app["bridge"]
    websockets: set = request.app["websockets"]
    websockets.add(ws)

    def on_event(name: str, fields: dict) -> None:
        _fire_and_forget(ws.send_json({"event": name, **fields}))

    bridge.add_event_listener(on_event)
    try:
        await ws.send_json({"event": "bridge_status", "upstream_connected": bridge.connected})

        async for msg in ws:
            if msg.type == WSMsgType.TEXT:
                await _handle_message(ws, bridge, msg.data)
            elif msg.type == WSMsgType.ERROR:
                logger.warning("websocket closed with exception %s", ws.exception())
    finally:
        bridge.remove_event_listener(on_event)
        websockets.discard(ws)

    return ws


def _fire_and_forget(coro) -> None:
    # Sends to a websocket that's mid-close routinely raise -- log at debug
    # rather than let asyncio print "exception never retrieved" warnings.
    task = asyncio.create_task(coro)

    def _log_if_failed(t: asyncio.Task) -> None:
        if not t.cancelled() and t.exception() is not None:
            logger.debug("websocket send failed: %s", t.exception())

    task.add_done_callback(_log_if_failed)


async def _handle_message(ws: web.WebSocketResponse, bridge: AntScopeZBridge, raw: str) -> None:
    try:
        msg = json.loads(raw)
    except json.JSONDecodeError:
        await ws.send_json({"ok": False, "error": "invalid JSON"})
        return

    client_id = msg.get("id")
    cmd = msg.get("cmd")
    if cmd not in ALLOWED_COMMANDS:
        await ws.send_json({"id": client_id, "ok": False, "error": f"unknown or unsupported command: {cmd!r}"})
        return

    params = {k: v for k, v in msg.items() if k not in ("cmd", "id")}
    try:
        reply = await bridge.request(cmd, **params)
    except (ConnectionError, TimeoutError) as exc:
        await ws.send_json({"id": client_id, "ok": False, "error": str(exc)})
        return
    except RuntimeError as exc:
        await ws.send_json({"id": client_id, "ok": False, "error": str(exc)})
        return

    reply["id"] = client_id
    await ws.send_json(reply)


async def on_startup(app: web.Application) -> None:
    bridge: AntScopeZBridge = app["bridge"]
    websockets: set = app["websockets"]

    def broadcast_status(connected: bool) -> None:
        for ws in list(websockets):
            _fire_and_forget(ws.send_json({"event": "bridge_status", "upstream_connected": connected}))

    bridge.add_status_listener(broadcast_status)
    bridge.start()


async def on_cleanup(app: web.Application) -> None:
    bridge: AntScopeZBridge = app["bridge"]
    await bridge.stop()
    for ws in list(app["websockets"]):
        await ws.close()


async def index_handler(request: web.Request) -> web.FileResponse:
    return web.FileResponse(STATIC_DIR / "index.html")


def build_app(antscopez_host: str, antscopez_port: int) -> web.Application:
    app = web.Application()
    app["bridge"] = AntScopeZBridge(antscopez_host, antscopez_port)
    app["websockets"] = set()

    app.on_startup.append(on_startup)
    app.on_cleanup.append(on_cleanup)

    app.router.add_get("/ws", ws_handler)
    # aiohttp's static mount doesn't serve index.html for "/" on its own
    # (unlike e.g. nginx), so that exact path needs an explicit route.
    app.router.add_get("/", index_handler)
    app.router.add_static("/", STATIC_DIR, show_index=False, name="static")

    return app


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--antscopez-host", default="127.0.0.1",
                         help="AntScopeZ Remote API host (default: 127.0.0.1 -- it only ever binds loopback)")
    parser.add_argument("--antscopez-port", type=int, default=7443,
                         help="AntScopeZ Remote API port (default: 7443, matches AntScopeZ's own default)")
    parser.add_argument("--listen-host", default="0.0.0.0",
                         help="address this bridge's web UI listens on (default: 0.0.0.0, so your phone can reach it)")
    parser.add_argument("--listen-port", type=int, default=8000,
                         help="port this bridge's web UI listens on (default: 8000)")
    parser.add_argument("-v", "--verbose", action="store_true", help="debug logging")
    args = parser.parse_args()

    logging.basicConfig(level=logging.DEBUG if args.verbose else logging.INFO,
                         format="%(asctime)s %(levelname)s %(name)s: %(message)s")

    app = build_app(args.antscopez_host, args.antscopez_port)
    logger.info("serving web UI on http://%s:%d/ (upstream AntScopeZ at %s:%d)",
                args.listen_host, args.listen_port, args.antscopez_host, args.antscopez_port)
    web.run_app(app, host=args.listen_host, port=args.listen_port, print=None)


if __name__ == "__main__":
    main()