ai-agent/sidecar/pairing.py

"""
Pairing — how a hub gets the bearer token of a worker it has never met.

A **worker** is this sidecar running on another machine (the Orus MacBook) as
a launchd agent, with no SIDECAR_TOKEN of its own. Its installer writes a
one-time **pairing code** and prints the string to paste into the hub's
Settings → Workers:

    100.80.162.92:8790/K7QMX4PJ2R/gabriel@orus.insure
    └─ tailnet addr ─┘ └─ code ─┘ └─ the Claude login ─┘

The hub then calls ``POST /pair {code, hubName}`` (the only unauthenticated
write here). A matching code is **consumed**, a fresh long-lived bearer is
minted and returned with the worker's identity and its login, and from then on
every other endpoint wants that bearer. The direction never flips: the hub
pulls, the worker never calls it back, so the worker holds no hub URL and no
hub credential.

State lives in ``WORKER_STATE_DIR`` (``~/.config/ai-agent-worker``), mode 0600:
``pairing-code`` (deleted on use), ``token`` (the bearer), ``worker.json``
(stable id + the hub's name). Re-pairing needs a new code (``install-macos.sh
--pair``) and rotates the bearer, so an old hub loses access. ``/unpair``
drops the bearer.

Wrong codes are counted; after ``MAX_ATTEMPTS`` the code is burned — the
listener is tailnet-only, but a guessable door on a work laptop is still a
door.
"""

from __future__ import annotations

import hmac
import json
import os
import pathlib
import platform
import secrets
import socket
import sys
import threading
import time
import uuid
from typing import Callable

from fastapi import APIRouter, Header, HTTPException
from pydantic import BaseModel

VERSION = "1"
MAX_ATTEMPTS = 10
# No 0/O/1/I/L: the code is read off one screen and typed or pasted on another.
_ALPHABET = "ABCDEFGHJKMNPQRSTUVWXYZ23456789"


def state_dir() -> pathlib.Path:
    return pathlib.Path(os.path.expanduser(os.environ.get(
        "WORKER_STATE_DIR", "~/.config/ai-agent-worker")))


def _file(name: str) -> pathlib.Path:
    return state_dir() / name


def _write_private(p: pathlib.Path, text: str) -> None:
    p.parent.mkdir(parents=True, exist_ok=True)
    tmp = p.with_name(p.name + ".tmp")
    fd = os.open(tmp, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
    with os.fdopen(fd, "w") as f:
        f.write(text)
    os.replace(tmp, p)


def _read(name: str) -> str:
    try:
        return _file(name).read_text().strip()
    except OSError:
        return ""


def token() -> str:
    """The bearer a paired hub presents ('' when unpaired)."""
    return _read("token")


def _worker() -> dict:
    try:
        return json.loads(_read("worker.json") or "{}")
    except ValueError:
        return {}


def identity() -> dict:
    """Who this runner is. The id is minted once and survives re-pairing, so
    a hub that re-pairs recognises the machine instead of adding a twin."""
    w = _worker()
    if not w.get("workerId"):
        w["workerId"] = uuid.uuid4().hex[:12]
        try:
            _write_private(_file("worker.json"), json.dumps(w))
        except OSError:
            pass
    host = socket.gethostname().split(".")[0]
    return {"workerId": w["workerId"], "hostname": host,
            "name": w.get("name") or host,
            "platform": sys.platform, "machine": platform.machine(),
            "version": VERSION, "hubName": w.get("hubName")}


def new_code() -> str:
    """Write a fresh one-time pairing code (replacing any unused one)."""
    code = "".join(secrets.choice(_ALPHABET) for _ in range(10))
    _write_private(_file("pairing-code"), code)
    _attempts[0] = 0
    return code


_lock = threading.Lock()
_attempts = [0]


class PairBody(BaseModel):
    code: str
    hubName: str | None = None       # shown by the worker's own logs/health


def make_router(auth: Callable[[str | None], None],
                account_status: Callable[[], dict],
                info: Callable[[], dict]) -> APIRouter:
    r = APIRouter()

    @r.post("/pair")
    def pair(body: PairBody) -> dict:
        """Trade the one-time code for the worker's bearer + identity."""
        with _lock:
            want = _read("pairing-code")
            got = (body.code or "").strip().upper()
            if not want:
                raise HTTPException(409, {
                    "message": "no pairing code on this worker — run "
                               "`install-macos.sh --pair` on it for a new one",
                    "kind": "no_code"})
            if not hmac.compare_digest(got, want):
                _attempts[0] += 1
                if _attempts[0] >= MAX_ATTEMPTS:
                    _file("pairing-code").unlink(missing_ok=True)
                raise HTTPException(403, {"message": "wrong pairing code",
                                          "kind": "bad_code"})
            tok = secrets.token_urlsafe(32)
            _write_private(_file("token"), tok)
            _file("pairing-code").unlink(missing_ok=True)
            w = _worker()
            w.update({"hubName": (body.hubName or "")[:80] or None,
                      "pairedAt": int(time.time())})
            _write_private(_file("worker.json"), json.dumps(w))
        print(f"[worker] paired with {body.hubName or 'a hub'}", flush=True)
        return {**identity(), **info(), "token": tok,
                "account": account_status()}

    @r.post("/unpair")
    def unpair(authorization: str | None = Header(default=None)) -> dict:
        """Forget the hub: the bearer stops working at once."""
        auth(authorization)
        _file("token").unlink(missing_ok=True)
        print("[worker] unpaired", flush=True)
        return {"ok": True}

    return r


if __name__ == "__main__":
    # `python pairing.py new-code` — what the installer calls to print the
    # pairing string. `status` says whether a hub holds the bearer.
    cmd = sys.argv[1] if len(sys.argv) > 1 else "status"
    if cmd == "new-code":
        print(new_code())
    else:
        print(json.dumps({**identity(), "paired": bool(token()),
                          "pendingCode": bool(_read("pairing-code"))}))