bookmarks-api/fab_auth.py

"""Ed25519 client registration and request signing.

The host stores public keys only. Private keys are written where the user
says and passed into tooling — never next to the native host state.
"""
from __future__ import annotations

import base64
import hashlib
import json
import os
import re
import time
import uuid
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

from cryptography.exceptions import InvalidSignature
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey, Ed25519PublicKey
from cryptography.hazmat.primitives.serialization import (
    Encoding,
    NoEncryption,
    PrivateFormat,
    PublicFormat,
    load_pem_private_key,
    load_pem_public_key,
)

STATE_DIR = Path(os.environ.get("LOCALAPPDATA", str(Path.home()))) / "firefox-agent-bridge"
CLIENTS_PATH = STATE_DIR / "clients.json"
SKEW_SECONDS = 90
NAME_RE = re.compile(r"^[A-Za-z0-9._-]{1,64}$")


def _now() -> str:
    return datetime.now(timezone.utc).replace(microsecond=0).isoformat()


def load_clients() -> list[dict[str, Any]]:
    if not CLIENTS_PATH.exists():
        return []
    data = json.loads(CLIENTS_PATH.read_text(encoding="utf-8"))
    return list(data.get("clients") or [])


def save_clients(clients: list[dict[str, Any]]) -> None:
    STATE_DIR.mkdir(parents=True, exist_ok=True)
    CLIENTS_PATH.write_text(
        json.dumps({"clients": clients}, indent=2) + "\n",
        encoding="utf-8",
    )


def _b64(data: bytes) -> str:
    return base64.b64encode(data).decode("ascii")


def _unb64(text: str) -> bytes:
    return base64.b64decode(text.encode("ascii"))


def body_hash(raw: bytes) -> str:
    return hashlib.sha256(raw).hexdigest()


def canonical(
    client_id: str,
    timestamp: str,
    nonce: str,
    http_method: str,
    path: str,
    raw_body: bytes,
) -> bytes:
    return "\n".join(
        [
            "v1",
            client_id,
            timestamp,
            nonce,
            http_method.upper(),
            path,
            body_hash(raw_body),
        ]
    ).encode("utf-8")


def register_client(name: str, key_path: Path) -> dict[str, str]:
    if not NAME_RE.match(name):
        raise ValueError("name must be 1-64 chars of A-Za-z0-9._-")
    key_path = key_path.expanduser().resolve()
    if STATE_DIR in key_path.parents or key_path.parent == STATE_DIR:
        raise ValueError(
            f"refusing to write a private key under {STATE_DIR} — pick a path you will give to tooling"
        )
    clients = load_clients()
    if any(c.get("name") == name for c in clients):
        raise ValueError(f"client already registered: {name}")
    private = Ed25519PrivateKey.generate()
    public = private.public_key()
    client_id = uuid.uuid4().hex
    pub_pem = public.public_bytes(Encoding.PEM, PublicFormat.SubjectPublicKeyInfo).decode("ascii")
    priv_pem = private.private_bytes(Encoding.PEM, PrivateFormat.PKCS8, NoEncryption()).decode("ascii")
    clients.append(
        {
            "id": client_id,
            "name": name,
            "public_key_pem": pub_pem,
            "created": _now(),
        }
    )
    save_clients(clients)
    key_path.parent.mkdir(parents=True, exist_ok=True)
    if key_path.exists():
        raise ValueError(f"key file already exists: {key_path}")
    bundle = {
        "id": client_id,
        "name": name,
        "private_key_pem": priv_pem,
        "public_key_pem": pub_pem,
    }
    key_path.write_text(json.dumps(bundle, indent=2) + "\n", encoding="utf-8")
    try:
        os.chmod(key_path, 0o600)
    except OSError:
        pass
    return {"id": client_id, "name": name, "key_file": str(key_path)}


def revoke_client(name_or_id: str) -> str:
    clients = load_clients()
    kept = [c for c in clients if c.get("id") != name_or_id and c.get("name") != name_or_id]
    if len(kept) == len(clients):
        raise ValueError(f"no client {name_or_id!r}")
    save_clients(kept)
    return name_or_id


def load_key_bundle(path: Path) -> dict[str, str]:
    data = json.loads(path.expanduser().read_text(encoding="utf-8"))
    if not data.get("id") or not data.get("private_key_pem"):
        raise ValueError("key file must contain id and private_key_pem")
    return data


def sign_headers(
    bundle: dict[str, str],
    http_method: str,
    path: str,
    raw_body: bytes,
) -> dict[str, str]:
    private = load_pem_private_key(bundle["private_key_pem"].encode("ascii"), password=None)
    if not isinstance(private, Ed25519PrivateKey):
        raise ValueError("key file is not an Ed25519 private key")
    timestamp = str(int(time.time()))
    nonce = uuid.uuid4().hex
    message = canonical(bundle["id"], timestamp, nonce, http_method, path, raw_body)
    signature = _b64(private.sign(message))
    return {
        "Authorization": f"FAB-ED25519 id={bundle['id']}",
        "X-FAB-Timestamp": timestamp,
        "X-FAB-Nonce": nonce,
        "X-FAB-Signature": signature,
    }


class ReplayCache:
    def __init__(self, limit: int = 2048) -> None:
        self.limit = limit
        self._seen: dict[str, float] = {}

    def accept(self, nonce: str, now: float) -> bool:
        cutoff = now - SKEW_SECONDS
        stale = [key for key, ts in self._seen.items() if ts < cutoff]
        for key in stale:
            del self._seen[key]
        if nonce in self._seen:
            return False
        self._seen[nonce] = now
        if len(self._seen) > self.limit:
            oldest = sorted(self._seen, key=self._seen.get)[: len(self._seen) - self.limit]
            for key in oldest:
                del self._seen[key]
        return True


def verify_request(
    headers: dict[str, str],
    http_method: str,
    path: str,
    raw_body: bytes,
    replay: ReplayCache,
) -> dict[str, Any]:
    auth = headers.get("Authorization") or headers.get("authorization") or ""
    match = re.fullmatch(r"FAB-ED25519 id=([0-9a-f]{32})", auth.strip())
    if not match:
        raise PermissionError("signed FAB-ED25519 client required")
    client_id = match.group(1)
    timestamp = (headers.get("X-FAB-Timestamp") or headers.get("x-fab-timestamp") or "").strip()
    nonce = (headers.get("X-FAB-Nonce") or headers.get("x-fab-nonce") or "").strip()
    signature = (headers.get("X-FAB-Signature") or headers.get("x-fab-signature") or "").strip()
    if not timestamp.isdigit() or not nonce or not signature:
        raise PermissionError("missing signature headers")
    now = time.time()
    if abs(now - int(timestamp)) > SKEW_SECONDS:
        raise PermissionError("timestamp outside allowed skew")
    if not replay.accept(nonce, now):
        raise PermissionError("replayed nonce")
    client = next((c for c in load_clients() if c.get("id") == client_id), None)
    if not client:
        raise PermissionError("unknown client id")
    public = load_pem_public_key(client["public_key_pem"].encode("ascii"))
    if not isinstance(public, Ed25519PublicKey):
        raise PermissionError("stored key is not Ed25519")
    message = canonical(client_id, timestamp, nonce, http_method, path, raw_body)
    try:
        public.verify(_unb64(signature), message)
    except InvalidSignature as exc:
        raise PermissionError("bad signature") from exc
    return client