distinct_agent/servers.py

"""Per-server credentials, pins, operator grants, and the admission gate.

This module is the agent's half of the handshake and it is the security
boundary.  The server filters what a user may select; that is a convenience and
it fails the moment the server is compromised or hostile.  What actually holds
is :class:`AllowanceGate`, which refuses any job naming a model or tool outside
the set the operator approved for the server that sent it, whatever the server
says.  ``tests/test_handshake_security.py`` exercises that path with a server
that has been told nothing and offers work anyway.

Three sets, and the distinction between them is the multi-server story:

``declared``
    The machine-wide ceiling, stated once at boot.  Never leaves the host.
``offered``
    Per server.  What that server is even told exists, chosen by the operator
    when the server is added.  This is what goes in ``AgentCapabilities`` at
    pairing, so a second server cannot learn that the first was offered a larger
    model.
``granted``
    Per server.  ``offered`` intersected with the server's catalogue, then
    narrowed by an explicit operator approval.  Empty until that approval
    happens.  This is what is advertised, and what the gate enforces.

The invariant ``granted <= offered <= declared`` holds at every point and is
asserted rather than assumed.

Key material
------------

Credentials live in memory by default, because a restart is then a blunt but
complete revocation and the README already promises it.  Persistence is opt-in
via ``path=``.  On POSIX the directory is created ``0o700`` and the file
``0o600``, both at creation time rather than afterwards, so there is no window
where the secret is world-readable.  **On Windows the file inherits the user
profile ACL**: that keeps it from other user accounts on the machine and does
nothing at all against another process running as the same user.  Say so rather
than implying a protection that is not there.

Rotation replaces the secret atomically (write a sibling temporary file, then
``os.replace``) and resets the snapshot counter, because a counter is only
meaningful relative to the key that signs it.  Rotation does **not** reopen the
approval question: the operator approved a set of models and tools, not a key.
A changed catalogue reopens it; a changed key does not.
"""

from __future__ import annotations

import json
import os
import secrets
import stat
import tempfile
import threading
import time
from collections.abc import Callable, Iterable, Mapping, Sequence
from dataclasses import dataclass, field
from pathlib import Path

from distinct_protocol import JobSpec
from distinct_protocol.handshake import (
    NOTHING,
    SNAPSHOT_ALGORITHM,
    Allowance,
    HandshakeError,
    OperatorApproval,
    ServerCatalogue,
    SignedSnapshot,
    ed25519_public_key,
    generate_ed25519_key,
)

#: Refusal reasons that travel to the server and therefore reach a browser.
#: They are a closed vocabulary of ASCII tokens and never interpolate a model
#: id, tool name or any other worker-supplied string, so nothing from this path
#: can carry markup into the UI.  The detail an operator needs stays local, in
#: :attr:`Refusal.detail`, which is written to the agent console only.
REASON_NOT_APPROVED = "not_approved"
REASON_MODEL_NOT_ALLOWED = "model_not_allowlisted"
REASON_TOOL_NOT_ALLOWED = "tool_not_allowlisted"
REASON_REVOKED = "revoked_by_operator"
REASON_UNKNOWN_SERVER = "unknown_server"

#: Refusals raised by the agent-local deny list rather than the allowlist. They
#: share this vocabulary so there is exactly one set of tokens that may cross
#: the wire, and so a reader can see the whole set in one place.
REASON_DENIED_BY_AGENT = "denied_by_agent"
REASON_UNAUTHENTICATED = "agent_requires_sign_in"

WIRE_REASONS = frozenset(
    {
        REASON_NOT_APPROVED,
        REASON_MODEL_NOT_ALLOWED,
        REASON_TOOL_NOT_ALLOWED,
        REASON_REVOKED,
        REASON_UNKNOWN_SERVER,
        REASON_DENIED_BY_AGENT,
        REASON_UNAUTHENTICATED,
    }
)


class AllowanceError(RuntimeError):
    """A handshake step was attempted out of order or with bad input."""


class NotApproved(AllowanceError):
    """No operator approval exists for this server yet."""


class PinMismatch(AllowanceError):
    """The server presented a different identity from the one pinned."""


class UnknownServer(AllowanceError):
    """No record exists for this server id."""


@dataclass(frozen=True)
class Refusal:
    """Why a job was refused.  ``reason`` is safe to send; ``detail`` is not."""

    reason: str
    detail: str

    def __post_init__(self) -> None:
        if self.reason not in WIRE_REASONS:
            raise AllowanceError(f"refusal reason {self.reason!r} is not in the closed set")


@dataclass(frozen=True)
class RevocationOutcome:
    """What a revocation removed and which local jobs it invalidates."""

    server_id: str
    removed: Allowance
    remaining: Allowance
    #: Ids of jobs the caller must now stop.  Computed by
    #: :meth:`ServerRegistry.revoke` from the jobs it was handed, so the sweep
    #: and the narrowing commit under the same lock and no job can slip between
    #: them.
    invalidated_job_ids: tuple[str, ...] = ()


@dataclass
class ServerRecord:
    """Everything the agent knows about one coordinator.

    ``secret`` is a ``bytearray`` so it can be overwritten on rotation.  Python
    gives no guarantee the bytes are gone from the heap after that, and this
    docstring is the honest version of that claim: overwriting removes the value
    from this object, not necessarily from the process.
    """

    server_id: str
    url: str
    agent_id: str
    pin: str
    offered: Allowance
    secret: bytearray = field(default_factory=bytearray, repr=False)
    #: The agent's Ed25519 signing key for this server's snapshots. **One
    #: keypair per server, not one per agent.** A single shared public key
    #: would be a stable identifier that two servers could compare to discover
    #: they are talking to the same machine, which is exactly the correlation
    #: the offer set and the per-server counter already remove.
    identity_key: bytearray = field(default_factory=bytearray, repr=False)
    catalogue: ServerCatalogue | None = None
    granted: Allowance = NOTHING
    approval: OperatorApproval | None = None
    counter: int = 0
    revoked: bool = False

    def __repr__(self) -> str:  # pragma: no cover - defensive, not behavioural
        return (
            f"ServerRecord(server_id={self.server_id!r}, agent_id={self.agent_id!r}, "
            f"granted={self.granted!r}, revoked={self.revoked!r})"
        )


class ServerRegistry:
    """Thread-safe store of per-server records, grants and counters."""

    def __init__(
        self,
        declared: Allowance,
        *,
        path: str | os.PathLike[str] | None = None,
        clock: Callable[[], float] = time.time,
    ) -> None:
        if not isinstance(declared, Allowance):
            raise AllowanceError("declared allowance must be an Allowance")
        self.declared = declared
        self._clock = clock
        self._path = Path(path) if path is not None else None
        self._lock = threading.RLock()
        self._records: dict[str, ServerRecord] = {}
        if self._path is not None and self._path.exists():
            self._load()

    # Registration -------------------------------------------------------

    def register(
        self,
        server_id: str,
        *,
        url: str,
        agent_id: str,
        pin: str,
        secret: bytes | bytearray,
        offer: Allowance,
        identity_key: bytes | bytearray | None = None,
        replace_existing: bool = False,
    ) -> ServerRecord:
        """Record a freshly paired server and the subset it was offered.

        ``offer`` is required rather than defaulted to the declared allowance.
        Defaulting it would mean that adding a second server silently disclosed
        everything the first one was offered, which is the leak this separation
        exists to prevent.

        ``replace_existing`` is for re-pairing with a server already on file.
        That is an ordinary thing to do: pairing codes are one-use and
        short-lived, so rejoining a server means a new code and new
        credentials, and refusing it would make the registry's memory a reason
        the worker could not start. It replaces the record outright, including
        dropping any approval attached to the old one, because an approval is
        bound to a catalogue digest and the new pairing has to re-take it. It
        is off by default so that the first registration of a duplicate id is
        still an error rather than a silent overwrite.
        """

        if not isinstance(offer, Allowance):
            raise AllowanceError("offer must be an Allowance")
        if not offer.issubset(self.declared):
            outside = offer.difference(self.declared)
            raise AllowanceError(
                "offer exceeds the declared allowance: "
                f"models {list(outside.models)}, tools {list(outside.tools)}"
            )
        if not isinstance(secret, (bytes, bytearray)) or len(secret) < 32:
            raise AllowanceError("server secret must contain at least 32 bytes")
        with self._lock:
            if server_id in self._records and not replace_existing:
                raise AllowanceError(f"server {server_id!r} is already registered")
            if identity_key is not None and len(identity_key) != 32:
                raise AllowanceError("an Ed25519 private key is exactly 32 bytes")
            record = ServerRecord(
                server_id=server_id,
                url=url,
                agent_id=agent_id,
                pin=pin,
                offered=offer,
                secret=bytearray(secret),
                # ``identity_key`` lets a caller that had to state its public
                # key *at pairing* (before this record could exist) register
                # the matching private key rather than minting a second one
                # the server never pinned.
                identity_key=bytearray(identity_key or generate_ed25519_key()),
            )
            self._records[server_id] = record
            self._persist_locked()
            return record

    def public_key(self, server_id: str) -> bytes:
        """The 32 raw bytes a server pins to check this agent's advertisements.

        Sent once, at pairing.  Everything after that is verifiable by anyone
        holding it, which is the property a shared secret cannot provide.
        """

        with self._lock:
            return ed25519_public_key(self._require_locked(server_id).identity_key)

    def verify_pin(self, server_id: str, presented: str) -> None:
        """Halt on a changed server identity rather than carrying on.

        Trust on first use: the first pin is taken on faith, which is weak
        against somebody already in position at that moment and strong against
        everything afterwards.  A mismatch is not a warning and not a retry; the
        caller must stop polling that server until an operator re-approves.
        """

        record = self.require(server_id)
        if not secrets.compare_digest(record.pin, str(presented)):
            raise PinMismatch(
                f"server {server_id!r} presented identity {str(presented)[:16]!r}..., "
                f"pinned {record.pin[:16]!r}.... Polling stops until an operator "
                "re-approves this server."
            )

    # Catalogue and approval ---------------------------------------------

    def record_catalogue(self, server_id: str, catalogue: ServerCatalogue) -> Allowance:
        """Store a pulled catalogue and return what the operator may approve.

        The choice set is ``catalogue.allowance & offered``.  A catalogue
        entry the operator never offered this server is not merely refused
        later, it is never shown, so it cannot be approved by accident.
        Recording a catalogue whose digest differs from the approved one
        **withdraws the existing grant**: consent was given for a specific set
        of permissions and a different set has not been consented to.
        """

        if not isinstance(catalogue, ServerCatalogue):
            raise AllowanceError("catalogue must be a ServerCatalogue")
        with self._lock:
            record = self._require_locked(server_id)
            if catalogue.server_id != server_id:
                raise AllowanceError("catalogue was issued by a different server")
            record.catalogue = catalogue
            if (
                record.approval is not None
                and record.approval.catalogue_digest != catalogue.digest()
            ):
                record.granted = NOTHING
                record.approval = None
            self._persist_locked()
            return catalogue.allowance.intersect(record.offered)

    def pending_choice(self, server_id: str) -> Allowance:
        """What is awaiting an operator decision for this server."""

        with self._lock:
            record = self._require_locked(server_id)
            if record.catalogue is None:
                return NOTHING
            return record.catalogue.allowance.intersect(record.offered)

    def approve(self, server_id: str, approval: OperatorApproval) -> Allowance:
        """Apply an operator's decision.  There is no auto-accept path.

        Every one of these checks has to pass, and each exists because skipping
        it would let something other than a human widen the grant:

        * a catalogue must have been pulled, so there is something to approve;
        * the approval's digest must match that catalogue, so consent given for
          one set of permissions cannot be carried over to another;
        * the approval must be for this server, so an approval for a different
          coordinator cannot be replayed here;
        * the approved set must sit inside the choice set, so approval can only
          narrow.
        """

        if not isinstance(approval, OperatorApproval):
            raise AllowanceError("approval must be an OperatorApproval")
        with self._lock:
            record = self._require_locked(server_id)
            if record.catalogue is None:
                raise AllowanceError(
                    f"no catalogue has been pulled from {server_id!r}; there is "
                    "nothing for an operator to approve"
                )
            if approval.server_id != server_id:
                raise AllowanceError(
                    "approval names a different server; approvals are per server "
                    "and are never shared between them"
                )
            expected = record.catalogue.digest()
            if not secrets.compare_digest(approval.catalogue_digest, expected):
                raise AllowanceError(
                    "approval is bound to a different catalogue. The server has "
                    "changed what it asks for since the operator was shown it, "
                    "so the earlier consent does not carry over."
                )
            choice = record.catalogue.allowance.intersect(record.offered)
            if not approval.allowance.issubset(choice):
                outside = approval.allowance.difference(choice)
                raise AllowanceError(
                    "approval exceeds what was offered to this server: "
                    f"models {list(outside.models)}, tools {list(outside.tools)}"
                )
            record.granted = approval.allowance
            record.approval = approval
            record.revoked = False
            self._persist_locked()
            return record.granted

    def granted(self, server_id: str) -> Allowance:
        """The approved set for one server, and only that server."""

        with self._lock:
            record = self._records.get(server_id)
            if record is None or record.revoked:
                return NOTHING
            return record.granted

    # Revocation ---------------------------------------------------------

    def revoke(
        self,
        server_id: str,
        *,
        models: Sequence[str] = (),
        tools: Sequence[str] = (),
        jobs: Iterable[JobSpec] = (),
    ) -> RevocationOutcome:
        """Narrow a grant and name the local work that must now stop.

        Withdrawing permission is present tense.  A job that is merely queued is
        cancelled outright: it has not started, and letting it through would run
        exactly the thing the operator just refused.  A job already running is
        also stopped, through the same cooperative cancellation the user-facing
        cancel uses, so the runner unwinds cleanly instead of being killed part
        way through a write.  Finishing "just this one" would be the wrong
        default: the operator revoked to stop it, not to stop the next one.

        The narrowing and the sweep happen under one lock, so no job can be
        admitted against the old grant after the new one is in force.
        """

        removed = Allowance(models=tuple(models), tools=tuple(tools))
        with self._lock:
            record = self._require_locked(server_id)
            before = record.granted
            record.granted = before.difference(removed)
            if not record.granted:
                record.revoked = True
            invalidated = tuple(
                job.id
                for job in jobs
                if isinstance(job, JobSpec) and _job_refusal(job, record.granted) is not None
            )
            self._persist_locked()
            return RevocationOutcome(
                server_id=server_id,
                removed=before.intersect(removed),
                remaining=record.granted,
                invalidated_job_ids=invalidated,
            )

    def forget(self, server_id: str) -> None:
        """Drop a server entirely and overwrite its secret."""

        with self._lock:
            record = self._records.pop(server_id, None)
            if record is not None:
                _zeroise(record.secret)
                _zeroise(record.identity_key)
            self._persist_locked()

    def rotate_secret(self, server_id: str, secret: bytes | bytearray) -> None:
        """Replace the transport secret; the operator's approval is untouched.

        This is the symmetric channel key, rotated by the server. The agent's
        Ed25519 identity is a different thing and is deliberately not touched
        here: see :meth:`rotate_identity`.
        """

        if not isinstance(secret, (bytes, bytearray)) or len(secret) < 32:
            raise AllowanceError("server secret must contain at least 32 bytes")
        with self._lock:
            record = self._require_locked(server_id)
            _zeroise(record.secret)
            record.secret = bytearray(secret)
            self._persist_locked()

    def rotate_identity(self, server_id: str) -> bytes:
        """Replace the signing identity, and invalidate the server's pin.

        Rotating the transport secret is routine and changes nothing an operator
        agreed to. **Rotating the identity is not routine.** The server pinned
        the old public key, so every advertisement signed with the new one will
        fail verification until the operator re-pairs and the server pins the
        new key. That is correct rather than inconvenient: the pin *is* the
        identity, and an identity that could be swapped silently would prove
        nothing.

        The counter resets, because a counter only orders assertions made under
        one key; carrying it over would let a snapshot signed with the old key
        and one signed with the new key claim the same position.
        """

        with self._lock:
            record = self._require_locked(server_id)
            _zeroise(record.identity_key)
            record.identity_key = bytearray(generate_ed25519_key())
            record.counter = 0
            self._persist_locked()
            return ed25519_public_key(record.identity_key)

    # Snapshots ----------------------------------------------------------

    def sign_snapshot(self, server_id: str, snapshot) -> SignedSnapshot:
        """Advertise to one server exactly what it was granted, and sign it.

        The counter increments once per snapshot per server.  Two servers never
        see the same counter value for the same state, and neither can infer the
        other's poll rate from it.
        """

        with self._lock:
            record = self._require_locked(server_id)
            record.counter += 1
            envelope = SignedSnapshot(
                snapshot=snapshot,
                advertised=NOTHING if record.revoked else record.granted,
                counter=record.counter,
                algorithm=SNAPSHOT_ALGORITHM,
                # Signed with the agent's own private key, which the server
                # never holds. The server can check the advertisement and
                # cannot manufacture one.
            ).signed(bytes(record.identity_key))
            self._persist_locked()
            return envelope

    def secret_for(self, server_id: str) -> bytes:
        with self._lock:
            return bytes(self._require_locked(server_id).secret)

    def require(self, server_id: str) -> ServerRecord:
        with self._lock:
            return self._require_locked(server_id)

    def server_ids(self) -> tuple[str, ...]:
        with self._lock:
            return tuple(sorted(self._records))

    def _require_locked(self, server_id: str) -> ServerRecord:
        record = self._records.get(server_id)
        if record is None:
            raise UnknownServer(f"no record for server {server_id!r}")
        return record

    # Persistence --------------------------------------------------------

    def _persist_locked(self) -> None:
        if self._path is None:
            return
        payload = {
            "declared": self.declared.to_dict(),
            "servers": [
                {
                    "server_id": record.server_id,
                    "url": record.url,
                    "agent_id": record.agent_id,
                    "pin": record.pin,
                    "offered": record.offered.to_dict(),
                    "granted": record.granted.to_dict(),
                    "approval": None if record.approval is None else record.approval.to_dict(),
                    "catalogue": None if record.catalogue is None else record.catalogue.to_dict(),
                    "secret": bytes(record.secret).hex(),
                    "identity_key": bytes(record.identity_key).hex(),
                    "counter": record.counter,
                    "revoked": record.revoked,
                }
                for record in (self._records[key] for key in sorted(self._records))
            ],
        }
        _write_private(self._path, json.dumps(payload, sort_keys=True, indent=2))

    def _load(self) -> None:
        assert self._path is not None
        raw = json.loads(self._path.read_text(encoding="utf-8"))
        if not isinstance(raw, Mapping):
            raise AllowanceError("server store is not a JSON object")
        for item in raw.get("servers") or ():
            if not isinstance(item, Mapping):
                raise AllowanceError("server store contains a malformed record")
            approval = item.get("approval")
            catalogue = item.get("catalogue")
            record = ServerRecord(
                server_id=str(item.get("server_id", "")),
                url=str(item.get("url", "")),
                agent_id=str(item.get("agent_id", "")),
                pin=str(item.get("pin", "")),
                offered=Allowance.from_dict(item.get("offered") or {}),
                secret=bytearray(bytes.fromhex(str(item.get("secret", "")))),
                identity_key=bytearray(bytes.fromhex(str(item.get("identity_key", "")))),
                catalogue=None if catalogue is None else ServerCatalogue.from_dict(catalogue),
                granted=Allowance.from_dict(item.get("granted") or {}),
                approval=None if approval is None else OperatorApproval.from_dict(approval),
                counter=int(item.get("counter", 0)),
                revoked=bool(item.get("revoked", False)),
            )
            # A stored grant is still only as good as the invariant.  If the
            # declared allowance narrowed while the agent was stopped, the
            # stored grant is trimmed on load rather than trusted.
            record.offered = record.offered.intersect(self.declared)
            record.granted = record.granted.intersect(record.offered)
            if len(record.identity_key) != 32:
                # A store written before signing identities existed. Mint one so
                # the agent still starts, and reset the counter, but note that
                # the server pinned nothing for this pairing: its advertisements
                # will not verify until the operator re-pairs.
                record.identity_key = bytearray(generate_ed25519_key())
                record.counter = 0
            self._records[record.server_id] = record


class AllowanceGate:
    """The agent-side check.  One server, one grant, no server input.

    Held deliberately small and free of any dependency on the transport, so the
    thing that decides whether a job may run can be read in one sitting and
    tested without a network, a queue or a server.
    """

    def __init__(self, registry: ServerRegistry, server_id: str) -> None:
        self.registry = registry
        self.server_id = server_id

    def check(self, job: JobSpec) -> Refusal | None:
        """Return ``None`` to admit, or a :class:`Refusal` to refuse.

        Refusing rather than raising is deliberate: a refusal is a normal,
        expected outcome that has to be reported back to the server and shown to
        the user, not an error condition on the worker.
        """

        if not isinstance(job, JobSpec):
            raise TypeError("job must be a distinct_protocol.JobSpec")
        try:
            record = self.registry.require(self.server_id)
        except UnknownServer:
            return Refusal(
                REASON_UNKNOWN_SERVER,
                f"job {job.id} arrived for unregistered server {self.server_id!r}",
            )
        if record.revoked:
            return Refusal(
                REASON_REVOKED,
                f"server {self.server_id!r} was revoked by the operator",
            )
        granted = self.registry.granted(self.server_id)
        if not granted:
            return Refusal(
                REASON_NOT_APPROVED,
                f"no operator approval exists for server {self.server_id!r}; "
                "the agent advertises nothing and accepts nothing",
            )
        return _job_refusal(job, granted)

    def advertised(self) -> Allowance:
        return self.registry.granted(self.server_id)


def _job_refusal(job: JobSpec, granted: Allowance) -> Refusal | None:
    """The one place a job is measured against a grant.

    Shared by the gate and by revocation so that "would this job be admitted?"
    and "does this running job still qualify?" can never answer differently.
    """

    if not granted.allows_model(job.model_id):
        return Refusal(
            REASON_MODEL_NOT_ALLOWED,
            f"model {job.model_id!r} is not in the approved set for this server",
        )
    for selection in job.allowed_tools:
        ref = f"{selection.id}@{selection.version}"
        if not granted.allows_tool(ref):
            return Refusal(
                REASON_TOOL_NOT_ALLOWED,
                f"tool {ref!r} is not in the approved set for this server",
            )
    return None


def _zeroise(buffer: bytearray) -> None:
    """Overwrite a secret in place.

    This clears the value from this object.  It does not promise the bytes are
    gone from the process: Python may have copied them, and there is no portable
    way to find or clear those copies.
    """

    for index in range(len(buffer)):
        buffer[index] = 0
    del buffer[:]


def _write_private(path: Path, text: str) -> None:
    """Write a secret-bearing file, never leaving it readable in between.

    The temporary file is created in the destination directory with the final
    permissions already set, then renamed over the target.  A reader either sees
    the old complete file or the new complete file, and never a partial one or a
    world-readable one.
    """

    directory = path.parent
    directory.mkdir(parents=True, exist_ok=True)
    if os.name != "nt":
        os.chmod(directory, stat.S_IRWXU)
    handle, temporary = tempfile.mkstemp(dir=str(directory), prefix=".servers-", suffix=".tmp")
    try:
        if os.name != "nt":
            os.fchmod(handle, stat.S_IRUSR | stat.S_IWUSR)
        with os.fdopen(handle, "w", encoding="utf-8") as stream:
            stream.write(text)
            stream.flush()
            os.fsync(stream.fileno())
        os.replace(temporary, path)
    except BaseException:
        os.unlink(temporary)
        raise


def storage_warning() -> str:
    """The line an operator must see before credentials touch a disk."""

    if os.name == "nt":
        return (
            "Stored agent credentials inherit this user profile's file "
            "permissions. That keeps them from other user accounts on this "
            "machine and does nothing against another process running as you. "
            "Without --remember-servers the credentials stay in memory and a "
            "restart revokes every pairing."
        )
    return (
        "Stored agent credentials are written 0600 in a 0700 directory, "
        "readable only by this user account. Without --remember-servers the "
        "credentials stay in memory and a restart revokes every pairing."
    )


__all__ = [
    "REASON_DENIED_BY_AGENT",
    "REASON_MODEL_NOT_ALLOWED",
    "REASON_NOT_APPROVED",
    "REASON_REVOKED",
    "REASON_TOOL_NOT_ALLOWED",
    "REASON_UNAUTHENTICATED",
    "REASON_UNKNOWN_SERVER",
    "WIRE_REASONS",
    "AllowanceError",
    "AllowanceGate",
    "HandshakeError",
    "NotApproved",
    "PinMismatch",
    "Refusal",
    "RevocationOutcome",
    "ServerRecord",
    "ServerRegistry",
    "UnknownServer",
    "storage_warning",
]