strix/tools/strix_serve.py

#!/usr/bin/env python3
"""HTTP server for a STRIX connected over USB, with the same routes as the device's own WiFi
HTTP API (arduino/strix/API.md), so clients work with either.

The USB port can only be opened by one program, so this server also follows the per-second
GCP-format trials itself: it polls the device, keeps the last hour in memory for /gcp, adds a
host timestamp to every record (`time`, Unix time at the start of the second) and can append
them to a JSON-lines file (--gcp-log). It reconnects after USB dropouts; requests meanwhile
get status 503.

With --set-token-file, /set requires the header `Authorization: Bearer <token>`, so the
server can be exposed publicly without letting anyone change settings.

/openapi.json serves the OpenAPI description (strix_openapi.json) and /docs a Swagger UI page,
which loads its scripts from a CDN.
"""

import argparse
import hmac
import json
import math
import threading
import time
from collections import deque
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
from urllib.parse import parse_qs, urlparse

import serial

from strix import MAX_ENTROPY_REQUEST, Device, find_port

MAX_HTTP_ENTROPY = 1 << 20
GCP_KEPT = 3600       # seconds of trials served by /gcp
GCP_POLL_S = 0.5      # also bounds the timestamp error
RECONNECT_S = 2
OPENAPI = (Path(__file__).parent / "strix_openapi.json").read_bytes()
SWAGGER_UI = b"""<!DOCTYPE html>
<html>
<head>
<meta charset="utf-8">
<title>STRIX HTTP API</title>
<link rel="stylesheet" href="https://unpkg.com/swagger-ui-dist@5/swagger-ui.css">
</head>
<body>
<div id="swagger-ui"></div>
<script src="https://unpkg.com/swagger-ui-dist@5/swagger-ui-bundle.js"></script>
<script>SwaggerUIBundle({url: "openapi.json", dom_id: "#swagger-ui"});</script>
</body>
</html>
"""
PASSTHROUGH = {"/id": "id", "/health": "health", "/stats": "stats", "/amp": "amp", "/wifi": "wifi",
               "/config": "get"}


class Unavailable(Exception):
    pass


class Strix:
    """The USB connection, shared by request threads and the GCP poller under one lock."""

    def __init__(self, port: str | None, gcp_log: str | None):
        self.port = port
        self.lock = threading.Lock()
        self.device = None
        self.records = deque(maxlen=GCP_KEPT)
        self.last_second = None
        self.reconnected = False
        self.offset = None  # host time at the end of second 0
        self.log = open(gcp_log, "a") if gcp_log else None

    def _connected(self) -> Device:
        if self.device is None:
            try:
                device = Device(self.port or find_port())
                # A client that exited without `watch off` leaves the device pushing lines that
                # would be taken for replies.
                device.send("watch off")
                while device.read_line() != '{"ok":true}':
                    pass
            except (serial.SerialException, OSError, TimeoutError) as error:
                raise Unavailable(str(error))
            self.device = device
        return self.device

    def call(self, action):
        """Runs action(device) under the lock; drops the connection on USB errors."""
        with self.lock:
            try:
                return action(self._connected())
            except (serial.SerialException, OSError, TimeoutError) as error:
                if self.device:
                    self.device.serial.close()
                self.device = None
                raise Unavailable(str(error))

    def command(self, line: str) -> str:
        def run(device):
            device.send(line)
            return device.read_line()
        return self.call(run)

    def entropy(self, count: int) -> bytes:
        def run(device):
            chunks = []
            while count > sum(map(len, chunks)):
                chunks.append(device.entropy(min(count - sum(map(len, chunks)), MAX_ENTROPY_REQUEST)))
            return b"".join(chunks)
        return self.call(run)

    def poll_gcp(self) -> None:
        reply = json.loads(self.command("gcp" if self.last_second is None else f"gcp {self.last_second}"))
        records = [reply] if "second" in reply else reply.get("records", [])
        if not records and self.reconnected:
            # A board reset restarts its second count, so nothing would ever follow last_second.
            latest = json.loads(self.command("gcp"))
            if latest.get("second", self.last_second) < self.last_second:
                self.last_second = None
            self.reconnected = False
        if not records:
            return
        newest = records[-1]["second"]
        now = time.time()
        # The newest record ended at most one poll interval ago; re-anchor after a device restart.
        if self.offset is None or abs(now - (self.offset + newest)) > 2:
            self.offset = math.floor(now) - newest
        for record in records:
            record = {"time": self.offset + record["second"] - 1, **record}
            self.records.append(record)
            if self.log:
                self.log.write(json.dumps(record) + "\n")
                self.log.flush()
        self.last_second = newest
        self.reconnected = False

    def run_poller(self) -> None:
        while True:
            try:
                self.poll_gcp()
                time.sleep(GCP_POLL_S)
            except Unavailable:
                self.reconnected = self.last_second is not None
                time.sleep(RECONNECT_S)
            except (ValueError, KeyError):
                time.sleep(GCP_POLL_S)  # no second completed yet, or a garbled line

    def gcp(self, since: str) -> dict:
        records = list(self.records)
        if not records:
            return {"error": "no second completed yet"}
        if not since:
            return records[-1]
        after = int(since)
        return {"records": [r for r in records if r["second"] > after]}


class Handler(BaseHTTPRequestHandler):
    strix: Strix
    set_token: str | None = None

    def may_set(self) -> bool:
        if self.set_token is None:
            return True
        given = self.headers.get("Authorization", "")
        return hmac.compare_digest(given.encode(), f"Bearer {self.set_token}".encode())

    def send(self, status: int, body: bytes, kind: str = "application/json") -> None:
        self.send_response(status)
        self.send_header("Content-Type", kind)
        self.send_header("Content-Length", str(len(body)))
        self.send_header("Access-Control-Allow-Origin", "*")
        self.end_headers()
        self.wfile.write(body)

    def send_json(self, value, status: int = 200) -> None:
        self.send(status, (json.dumps(value) + "\n").encode())

    def do_GET(self) -> None:
        url = urlparse(self.path)
        query = {key: values[0] for key, values in parse_qs(url.query).items()}
        try:
            if url.path == "/entropy" or (url.path == "/" and "length" in query):
                self.serve_entropy(query.get("bytes", query.get("length", "")), query.get("format", "hex"))
            elif url.path in PASSTHROUGH:
                self.send(200, (self.strix.command(PASSTHROUGH[url.path]) + "\n").encode())
            elif url.path == "/set" and not self.may_set():
                self.send_json({"error": "unauthorized"}, 401)
            elif url.path == "/set":
                line = f"set {query.get('key', '')} {query.get('value', '')}"
                self.send(200, (self.strix.command(line) + "\n").encode())
            elif url.path == "/gcp":
                self.send_json(self.strix.gcp(query.get("since", "")))
            elif url.path == "/openapi.json":
                self.send(200, OPENAPI)
            elif url.path == "/docs":
                self.send(200, SWAGGER_UI, "text/html; charset=utf-8")
            else:
                self.send_json({"error": "not found"}, 404)
        except Unavailable as error:
            self.send_json({"error": f"device unavailable: {error}"}, 503)
        except ValueError:
            self.send_json({"error": "invalid parameter"})

    def serve_entropy(self, count_text: str, form: str) -> None:
        count = int(count_text)
        if not 1 <= count <= MAX_HTTP_ENTROPY or form not in ("hex", "bin"):
            self.send_json({"error": f"usage: bytes=<1-{MAX_HTTP_ENTROPY}>&format=hex|bin"})
            return
        try:
            data = self.strix.entropy(count)
        except RuntimeError as error:  # the device's error reply, e.g. an unhealthy source
            self.send(200, (str(error) + "\n").encode())
            return
        if form == "bin":
            self.send(200, data, "application/octet-stream")
        else:
            self.send(200, (data.hex().upper() + "\n").encode(), "text/plain")

    def log_message(self, format, *args) -> None:
        pass


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--listen", default="127.0.0.1:8080", help="address:port (default 127.0.0.1:8080)")
    parser.add_argument("--port", help="serial port (default: the STRIX board, or STRIX_PORT)")
    parser.add_argument("--gcp-log", help="append every per-second trial to this JSON-lines file")
    parser.add_argument("--set-token-file", help="require this file's token as a bearer token for /set")
    args = parser.parse_args()

    if args.set_token_file:
        with open(args.set_token_file) as token:
            Handler.set_token = token.read().strip()
        if not Handler.set_token:
            parser.error("the set token file is empty")

    Handler.strix = Strix(args.port, args.gcp_log)
    threading.Thread(target=Handler.strix.run_poller, daemon=True).start()
    host, port = args.listen.rsplit(":", 1)
    server = ThreadingHTTPServer((host, int(port)), Handler)
    print(f"serving on http://{host}:{port}", flush=True)
    server.serve_forever()


if __name__ == "__main__":
    main()