Spaces:

User1342
/

distinct

Sleeping

distinct / distinct_protocol /handshake.py
initial commit
2aa8b3a about 2 months ago
30.8 kB
"""Wire objects for the agent registration and pairing handshake.
The handshake exists because an agent runs on somebody's own computer and the
operator of that computer, not the server, decides what it will do.  Five steps,
in this order, and the order is the point:
1. **Declare.** On boot the agent states an :class:`Allowance`: the models it is
   willing to run and the tools it is willing to execute.  This is the ceiling
   for the whole machine.
2. **Offer and register.** The operator enters a server address and chooses, per
   server, which subset of the declared allowance that server is even told
   about.  The agent then pairs.  A server never learns the machine-wide
   allowance, only its own offer set, so two servers cannot compare notes.
3. **Pull.** The agent fetches the server's :class:`ServerCatalogue`: everything
   that server might ask for.
4. **Approve.** The operator approves a subset.  The approval is bound to the
   digest of the exact catalogue that was shown, so a later catalogue cannot
   inherit an earlier consent.  There is no auto-accept.
5. **Advertise.** The approved :class:`Allowance` travels in every signed
   snapshot, and it is what the server filters user choices against.
Server-side filtering is a convenience.  The control is agent-side: the worker
refuses any job naming a model or tool outside its own approved set, whatever
the server sends.  See :mod:`distinct_agent.servers`.
Signatures are HMAC-SHA256 with **domain separation**, so a signature minted for
one purpose can never be presented as another.  Every signing call names its
domain; :data:`DOMAINS` is the closed set.
.. warning::
   HMAC is symmetric.  The server holds the same key, so a signature proves
   authenticity **to that server** and is not evidence a third party could
   check.  See ``PLAN.md`` section 4.8 for the recommendation and for what an
   asymmetric agent key would buy.
"""

from __future__ import annotations

import hashlib
import hmac
import re
from dataclasses import dataclass, replace
from typing import Any, Dict, Iterable, Mapping, Optional, Tuple

from .models import (
    AgentSnapshot,
    ProtocolError,
    _finite,
    _identifier,
    _text,
    canonical_json,
)

#: Bumped independently of ``PROTOCOL_VERSION``: the handshake can gain fields
#: without forcing a hard cut-over on the job wire.
HANDSHAKE_VERSION = "1.0"

#: Symmetric.  Used for the transport channel, where both parties already share
#: a per-agent secret and the question being answered is "did this exchange come
#: from the party I paired with", which a shared key answers perfectly well.
ALG_HMAC = "hmac-sha256"

#: Asymmetric.  Used for the detached snapshot signature, where the question is
#: different: **who said this**, answerable by someone who was not present.
ALG_ED25519 = "ed25519"

SIGNATURE_ALGORITHMS = frozenset({ALG_HMAC, ALG_ED25519})

#: The transport default, and the historical meaning of this name.
SIGNATURE_ALGORITHM = ALG_HMAC

#: The snapshot advertises which models and tools an operator approved, and that
#: is the field any dispute turns on.  HMAC proves only that *one of the two*
#: parties holding the key produced the message, so a server can forge an
#: agent's advertisement indistinguishably: integrity, but not evidence.  If an
#: operator ever has to show they did not approve something, or a user has to
#: show an agent claimed a capability it did not have, a shared key cannot
#: answer it.  Ed25519 gives third-party verifiability, which is the entire
#: point of having made the signature detached in the first place: a signature
#: that survives being unwrapped is worth little if only the two parties who
#: could each have forged it can check it.
SNAPSHOT_ALGORITHM = ALG_ED25519

#: Domain separation strings.  A signature is computed over
#: ``domain + "\n" + canonical_json(payload)``, so a snapshot signature can
#: never be replayed as a catalogue signature even with an identical body.
DOMAIN_CATALOGUE = "distinct/handshake/catalogue/1"
DOMAIN_PAIR_RESPONSE = "distinct/handshake/pair-response/1"
DOMAIN_RESPONSE = "distinct/handshake/response/1"
DOMAIN_SNAPSHOT = "distinct/handshake/snapshot/1"

DOMAINS = frozenset(
    {DOMAIN_CATALOGUE, DOMAIN_PAIR_RESPONSE, DOMAIN_RESPONSE, DOMAIN_SNAPSHOT}
)

#: Mirrors ``distinct_tools.core.ToolRef``: ``tool_id`` is a lowercase dotted
#: identifier, ``version`` a simple identifier.  The two regexes must stay in
#: step, and ``tests/test_handshake.py`` asserts that they do rather than
#: leaving it to a comment.  Allowlisting a tool by any other name would create
#: a second, drifting tool identity scheme.
TOOL_ID_PATTERN = r"[a-z][a-z0-9]*(?:[._-][a-z0-9]+)*"
TOOL_VERSION_PATTERN = r"[A-Za-z0-9][A-Za-z0-9._-]*"
_TOOL_REF_RE = re.compile(rf"^({TOOL_ID_PATTERN})@({TOOL_VERSION_PATTERN})$")

_CHALLENGE_RE = re.compile(r"^[A-Za-z0-9._~-]{16,128}$")
_DIGEST_RE = re.compile(r"^[0-9a-f]{64}$")

#: Bounds on how large an allowance may be.  They match
#: ``ControlPlaneLimits.max_models_per_agent`` and ``max_tools_per_agent`` so a
#: grant the agent considers valid cannot be one the server would refuse.
MAX_MODELS = 64
MAX_TOOLS = 64

#: A monotonic snapshot counter is a 63-bit integer.  At one poll per second it
#: would take longer than the age of the universe to wrap, so the ceiling is a
#: validation bound rather than a lifetime concern.
MAX_COUNTER = 2**63 - 1


class HandshakeError(ProtocolError):
    """Raised when a handshake value is malformed or fails verification."""


def parse_tool_ref(value: Any) -> Tuple[str, str]:
    """Split ``"id@version"`` into its parts, or refuse it.
    Versions are never resolved implicitly.  ``web.search`` without a version is
    not a tool reference, because "whichever version happens to be installed" is
    not something an operator can meaningfully approve.
    """

    if not isinstance(value, str):
        raise HandshakeError("tool reference must be text")
    match = _TOOL_REF_RE.fullmatch(value)
    if not match:
        raise HandshakeError(
            f"tool reference {value[:64]!r} must have the exact form id@version"
        )
    if len(match.group(1)) > 96 or len(match.group(2)) > 32:
        raise HandshakeError("tool reference is too long")
    return match.group(1), match.group(2)


def _unique_sorted(values: Iterable[Any], label: str, maximum: int) -> Tuple[str, ...]:
    items = list(values)
    if len(items) > maximum:
        raise HandshakeError(f"{label} exceeds {maximum} entries")
    seen = set()
    for item in items:
        if item in seen:
            raise HandshakeError(f"{label} contains a duplicate entry")
        seen.add(item)
    # Sorted so that two allowances holding the same members produce the same
    # digest.  An approval is bound to a digest, so an unstable order would make
    # a re-presented catalogue look like a new one.
    return tuple(sorted(seen))


@dataclass(frozen=True)
class Allowance:
    """A set of model ids and exact tool references.
    Empty is a legitimate and meaningful value: it is what an agent advertises
    to a server whose catalogue the operator has not approved.  An agent with an
    empty allowance is reachable and online, and will run nothing.
    """

    models: Tuple[str, ...] = ()
    tools: Tuple[str, ...] = ()

    def __post_init__(self) -> None:
        models = _unique_sorted(
            (_identifier(value, "model id") for value in self.models),
            "models",
            MAX_MODELS,
        )
        tools = []
        for value in self.tools:
            tool_id, version = parse_tool_ref(value)
            tools.append(f"{tool_id}@{version}")
        object.__setattr__(self, "models", models)
        object.__setattr__(self, "tools", _unique_sorted(tools, "tools", MAX_TOOLS))

    def __bool__(self) -> bool:
        return bool(self.models or self.tools)

    def allows_model(self, model_id: Any) -> bool:
        return isinstance(model_id, str) and model_id in self.models

    def allows_tool(self, ref: Any) -> bool:
        return isinstance(ref, str) and ref in self.tools

    def intersect(self, other: "Allowance") -> "Allowance":
        """The largest allowance permitted by both.
        Every narrowing in the handshake is an intersection, which is why a
        server cannot widen anything by sending more: adding members to one side
        of an intersection can only ever produce a subset of the other side.
        """

        if not isinstance(other, Allowance):
            raise HandshakeError("intersect requires another Allowance")
        return Allowance(
            models=tuple(set(self.models) & set(other.models)),
            tools=tuple(set(self.tools) & set(other.tools)),
        )

    def difference(self, other: "Allowance") -> "Allowance":
        """What ``self`` permits and ``other`` does not."""

        if not isinstance(other, Allowance):
            raise HandshakeError("difference requires another Allowance")
        return Allowance(
            models=tuple(set(self.models) - set(other.models)),
            tools=tuple(set(self.tools) - set(other.tools)),
        )

    def issubset(self, other: "Allowance") -> bool:
        if not isinstance(other, Allowance):
            raise HandshakeError("issubset requires another Allowance")
        return set(self.models).issubset(other.models) and set(self.tools).issubset(
            other.tools
        )

    def digest(self) -> str:
        """A stable identity for this exact set of permissions."""

        return hashlib.sha256(canonical_json(self.to_dict()).encode("utf-8")).hexdigest()

    def to_dict(self) -> Dict[str, Any]:
        return {"models": list(self.models), "tools": list(self.tools)}

    @classmethod
    def from_dict(cls, value: Mapping[str, Any]) -> "Allowance":
        if not isinstance(value, Mapping):
            raise HandshakeError("allowance must be a JSON object")
        return cls(
            models=tuple(value.get("models") or ()),
            tools=tuple(value.get("tools") or ()),
        )


#: The allowance an agent advertises before an operator has approved anything.
NOTHING = Allowance()


def _challenge(value: Any) -> str:
    if not isinstance(value, str) or not _CHALLENGE_RE.fullmatch(value):
        raise HandshakeError("challenge must be 16-128 URL-safe characters")
    return value


def _digest_value(value: Any, label: str) -> str:
    if not isinstance(value, str) or not _DIGEST_RE.fullmatch(value):
        raise HandshakeError(f"{label} must be 64 lowercase hexadecimal characters")
    return value


@dataclass(frozen=True)
class ServerCatalogue:
    """Everything one server may ask an agent to do, as that server states it.
    ``challenge`` is the agent's own nonce, echoed back.  It is what makes the
    catalogue fresh: a catalogue captured from an earlier exchange carries an
    earlier challenge and is refused before the operator is ever shown it.  A
    timestamp was the alternative and was rejected, because it would make
    freshness depend on clock agreement between two parties who do not yet trust
    each other, and the agent has no way to check the server's clock.
    """

    server_id: str
    challenge: str
    allowance: Allowance
    issued_at: float
    signature: str = ""

    def __post_init__(self) -> None:
        object.__setattr__(self, "server_id", _identifier(self.server_id, "server id"))
        object.__setattr__(self, "challenge", _challenge(self.challenge))
        if not isinstance(self.allowance, Allowance):
            object.__setattr__(self, "allowance", Allowance.from_dict(self.allowance))
        object.__setattr__(self, "issued_at", _finite(self.issued_at, "issued_at"))
        object.__setattr__(
            self, "signature", _text(self.signature, "signature", 128, allow_empty=True)
        )

    def signing_payload(self) -> Dict[str, Any]:
        return {
            "handshake_version": HANDSHAKE_VERSION,
            "alg": SIGNATURE_ALGORITHM,
            "server_id": self.server_id,
            "challenge": self.challenge,
            "allowance": self.allowance.to_dict(),
            "issued_at": self.issued_at,
        }

    def digest(self) -> str:
        """What an operator approval is bound to.
        The digest covers the allowance and the issuing server, and deliberately
        **not** the challenge or the timestamp.  Approving a catalogue is a
        decision about a set of permissions from a named server, so re-fetching
        an unchanged catalogue must not invalidate the consent the operator
        already gave.  Changing what the server asks for does.
        """

        return hashlib.sha256(
            canonical_json(
                {
                    "handshake_version": HANDSHAKE_VERSION,
                    "server_id": self.server_id,
                    "allowance": self.allowance.to_dict(),
                }
            ).encode("utf-8")
        ).hexdigest()

    def signed(self, secret: Any) -> "ServerCatalogue":
        return replace(
            self, signature=sign_payload(secret, DOMAIN_CATALOGUE, self.signing_payload())
        )

    def verify(self, secret: Any, *, challenge: Optional[str] = None) -> bool:
        """True only when the signature matches and the challenge is ours."""

        if challenge is not None and not hmac.compare_digest(
            self.challenge, _challenge(challenge)
        ):
            return False
        return verify_payload(
            secret, DOMAIN_CATALOGUE, self.signing_payload(), self.signature
        )

    def to_dict(self) -> Dict[str, Any]:
        payload = self.signing_payload()
        payload["signature"] = self.signature
        return payload

    @classmethod
    def from_dict(cls, value: Mapping[str, Any]) -> "ServerCatalogue":
        if not isinstance(value, Mapping):
            raise HandshakeError("catalogue must be a JSON object")
        version = value.get("handshake_version", HANDSHAKE_VERSION)
        if version != HANDSHAKE_VERSION:
            raise HandshakeError(f"unsupported handshake version {version!r}")
        algorithm = value.get("alg", SIGNATURE_ALGORITHM)
        if algorithm != SIGNATURE_ALGORITHM:
            raise HandshakeError(f"unsupported signature algorithm {algorithm!r}")
        return cls(
            server_id=value.get("server_id", ""),
            challenge=value.get("challenge", ""),
            allowance=Allowance.from_dict(value.get("allowance") or {}),
            issued_at=value.get("issued_at", 0),
            signature=value.get("signature", ""),
        )


@dataclass(frozen=True)
class OperatorApproval:
    """A record that a human said yes to a specific set of permissions.
    ``catalogue_digest`` binds the decision to the exact catalogue that was
    shown.  ``method`` records how consent was obtained and is written into the
    agent's console output, so an approval that arrived by some route nobody
    intended is visible rather than indistinguishable from a real one.
    """

    server_id: str
    catalogue_digest: str
    allowance: Allowance
    approved_at: float
    method: str = "operator-console"

    def __post_init__(self) -> None:
        object.__setattr__(self, "server_id", _identifier(self.server_id, "server id"))
        object.__setattr__(
            self,
            "catalogue_digest",
            _digest_value(self.catalogue_digest, "catalogue digest"),
        )
        if not isinstance(self.allowance, Allowance):
            object.__setattr__(self, "allowance", Allowance.from_dict(self.allowance))
        object.__setattr__(self, "approved_at", _finite(self.approved_at, "approved_at"))
        object.__setattr__(self, "method", _text(self.method, "approval method", 120))

    def to_dict(self) -> Dict[str, Any]:
        return {
            "handshake_version": HANDSHAKE_VERSION,
            "server_id": self.server_id,
            "catalogue_digest": self.catalogue_digest,
            "allowance": self.allowance.to_dict(),
            "approved_at": self.approved_at,
            "method": self.method,
        }

    @classmethod
    def from_dict(cls, value: Mapping[str, Any]) -> "OperatorApproval":
        if not isinstance(value, Mapping):
            raise HandshakeError("approval must be a JSON object")
        return cls(
            server_id=value.get("server_id", ""),
            catalogue_digest=value.get("catalogue_digest", ""),
            allowance=Allowance.from_dict(value.get("allowance") or {}),
            approved_at=value.get("approved_at", 0),
            method=value.get("method", "operator-console"),
        )


@dataclass(frozen=True)
class SignedSnapshot:
    """An :class:`AgentSnapshot` plus what the agent will accept, signed.
    The signature is **detached**: it covers the snapshot's own canonical form,
    not the request that carried it.  So the advertisement survives being
    unwrapped, stored, rendered into a table or quoted back in a dispute, and
    can still be checked.  A transit-only signature would have made the
    advertisement unverifiable the instant the request envelope was opened,
    which matters now that the snapshot is what says which models and tools this
    machine agreed to run.
    ``counter`` is a per-server monotonic sequence number.  A snapshot is an
    assertion about current state, so what a receiver needs to reject is an
    *older* one, and that is an ordering question rather than a freshness
    question.  A timestamp would have made ordering depend on two clocks
    agreeing; the counter needs no clock at all.
    """

    snapshot: AgentSnapshot
    advertised: Allowance
    counter: int
    signature: str = ""
    algorithm: str = SNAPSHOT_ALGORITHM

    def __post_init__(self) -> None:
        if self.algorithm not in SIGNATURE_ALGORITHMS:
            raise HandshakeError(f"unsupported signature algorithm {self.algorithm!r}")
        if not isinstance(self.snapshot, AgentSnapshot):
            raise HandshakeError("snapshot must be an AgentSnapshot")
        if not isinstance(self.advertised, Allowance):
            object.__setattr__(self, "advertised", Allowance.from_dict(self.advertised))
        if isinstance(self.counter, bool) or not isinstance(self.counter, int):
            raise HandshakeError("counter must be an integer")
        if not 1 <= self.counter <= MAX_COUNTER:
            raise HandshakeError("counter is outside the supported range")
        object.__setattr__(
            self, "signature", _text(self.signature, "signature", 128, allow_empty=True)
        )

    def signing_payload(self) -> Dict[str, Any]:
        return {
            "handshake_version": HANDSHAKE_VERSION,
            "alg": self.algorithm,
            "snapshot": self.snapshot.to_dict(),
            "advertised": self.advertised.to_dict(),
            "counter": self.counter,
        }

    def signed(self, key: Any) -> "SignedSnapshot":
        """Sign with the agent's private key.  ``key`` is 32 raw bytes."""

        return replace(
            self,
            signature=sign_payload(
                key, DOMAIN_SNAPSHOT, self.signing_payload(), algorithm=self.algorithm
            ),
        )

    def verify(self, key: Any, *, algorithm: Optional[str] = None) -> bool:
        """Verify against a **pinned** key, under a pinned algorithm.
        ``algorithm`` defaults to this envelope's own field only for the
        convenience of a caller that has already checked it.  A server holding a
        pinned Ed25519 public key should pass :data:`ALG_ED25519` explicitly, so
        that an envelope arriving with ``alg`` set to ``hmac-sha256`` is refused
        rather than verified under whichever scheme the sender chose.
        """

        expected = algorithm or self.algorithm
        if expected != self.algorithm:
            return False
        return verify_payload(
            key,
            DOMAIN_SNAPSHOT,
            self.signing_payload(),
            self.signature,
            algorithm=expected,
        )

    def to_dict(self) -> Dict[str, Any]:
        payload = self.signing_payload()
        payload["signature"] = self.signature
        return payload

    @classmethod
    def from_dict(cls, value: Mapping[str, Any]) -> "SignedSnapshot":
        if not isinstance(value, Mapping):
            raise HandshakeError("signed snapshot must be a JSON object")
        version = value.get("handshake_version", HANDSHAKE_VERSION)
        if version != HANDSHAKE_VERSION:
            raise HandshakeError(f"unsupported handshake version {version!r}")
        algorithm = value.get("alg", SNAPSHOT_ALGORITHM)
        if algorithm not in SIGNATURE_ALGORITHMS:
            raise HandshakeError(f"unsupported signature algorithm {algorithm!r}")
        raw_snapshot = value.get("snapshot")
        if not isinstance(raw_snapshot, Mapping):
            raise HandshakeError("signed snapshot omitted its snapshot")
        return cls(
            snapshot=AgentSnapshot.from_dict(raw_snapshot),
            advertised=Allowance.from_dict(value.get("advertised") or {}),
            counter=value.get("counter", 0),
            signature=value.get("signature", ""),
            algorithm=algorithm,
        )


@dataclass(frozen=True)
class ResponseAuthentication:
    """The proof a server attaches to everything the agent will act on.
    Under pull the agent dials an address a human typed in, so the direction
    that carries prompts and tool invocations is the one that most needs
    authenticating.  The bound value is the agent's **own request nonce**, which
    the agent generated moments earlier and has never reused.  That gives replay
    resistance with no clock agreement, no server-side cache and no state on
    either side: a captured response cannot be presented against any later
    request, because every later request carries a different nonce.
    """

    path: str
    request_nonce: str
    body_digest: str
    signature: str = ""

    def __post_init__(self) -> None:
        object.__setattr__(self, "path", _text(self.path, "path", 256))
        object.__setattr__(self, "request_nonce", _challenge(self.request_nonce))
        object.__setattr__(
            self, "body_digest", _digest_value(self.body_digest, "body digest")
        )
        object.__setattr__(
            self, "signature", _text(self.signature, "signature", 128, allow_empty=True)
        )

    def signing_payload(self) -> Dict[str, Any]:
        return {
            "handshake_version": HANDSHAKE_VERSION,
            "alg": SIGNATURE_ALGORITHM,
            "direction": "server-to-agent",
            "path": self.path,
            "request_nonce": self.request_nonce,
            "body_digest": self.body_digest,
        }

    def signed(self, secret: Any) -> "ResponseAuthentication":
        return replace(
            self, signature=sign_payload(secret, DOMAIN_RESPONSE, self.signing_payload())
        )

    def verify(self, secret: Any) -> bool:
        return verify_payload(
            secret, DOMAIN_RESPONSE, self.signing_payload(), self.signature
        )


#: The key a signed response carries its proof under.
RESPONSE_AUTH_KEY = "auth"


def sign_response(
    secret: Any,
    path: str,
    request_nonce: str,
    body: Mapping[str, Any],
) -> Dict[str, Any]:
    """Build the proof a server attaches to a response body.
    Lives in the protocol package because both sides need it and neither may
    import the other: the server mints these, the agent checks them, and a
    verifier tested only against its own construction proves nothing.  The
    digest is taken over the body *without* the proof key, which is the form
    :func:`verify_response` reconstructs.
    """

    authentication = ResponseAuthentication(
        path=path,
        request_nonce=request_nonce,
        body_digest=body_digest(canonical_json(dict(body))),
    ).signed(secret)
    return {
        "path": authentication.path,
        "request_nonce": authentication.request_nonce,
        "body_digest": authentication.body_digest,
        "signature": authentication.signature,
    }


def body_digest(body: Any) -> str:
    """The sha256 of a response body, in the one form both sides compute."""

    if isinstance(body, str):
        body = body.encode("utf-8")
    if not isinstance(body, (bytes, bytearray, memoryview)):
        raise HandshakeError("body must be bytes or text")
    return hashlib.sha256(bytes(body)).hexdigest()


def _key(secret: Any) -> bytes:
    if isinstance(secret, str):
        secret = secret.encode("utf-8")
    if isinstance(secret, (bytearray, memoryview)):
        secret = bytes(secret)
    if not isinstance(secret, bytes):
        raise HandshakeError("signing secret must be bytes or text")
    if len(secret) < 32:
        raise HandshakeError("signing secret must contain at least 32 bytes")
    return secret


def signing_input(domain: str, payload: Mapping[str, Any]) -> bytes:
    """The exact bytes signed, with the domain outside the JSON.
    Putting the domain outside means it cannot be confused with a payload field
    that an attacker controls.
    """

    if domain not in DOMAINS:
        raise HandshakeError(f"unknown signing domain {domain!r}")
    if not isinstance(payload, Mapping):
        raise HandshakeError("signing payload must be a mapping")
    return f"{domain}\n{canonical_json(dict(payload))}".encode("utf-8")


class SigningUnavailable(HandshakeError):
    """Asymmetric signing was requested and the backend is not installed."""


def _ed25519_module():
    """Import the Ed25519 primitives, or say exactly what is missing.
    Deliberately never falls back to HMAC.  A silent downgrade would return the
    scheme to one where the server can forge the agent's advertisement, which is
    the precise property this algorithm was chosen to remove, and it would do so
    invisibly.  Failing loudly is the only safe behaviour.
    """

    try:
        from cryptography.hazmat.primitives.asymmetric import ed25519
    except ImportError as exc:  # pragma: no cover - exercised only without the dep
        raise SigningUnavailable(
            "Ed25519 signing needs the 'cryptography' package. The detached "
            "snapshot signature will not fall back to HMAC, because a shared "
            "key cannot prove which party produced a message."
        ) from exc
    return ed25519


def generate_ed25519_key() -> bytes:
    """Return a new 32-byte raw Ed25519 private key."""

    ed25519 = _ed25519_module()
    return ed25519.Ed25519PrivateKey.generate().private_bytes_raw()


def ed25519_public_key(private_key: bytes | bytearray) -> bytes:
    """Derive the 32-byte raw public key that a server pins."""

    ed25519 = _ed25519_module()
    if not isinstance(private_key, (bytes, bytearray)) or len(private_key) != 32:
        raise HandshakeError("an Ed25519 private key is exactly 32 bytes")
    loaded = ed25519.Ed25519PrivateKey.from_private_bytes(bytes(private_key))
    return loaded.public_key().public_bytes_raw()


def sign_payload(
    key: Any,
    domain: str,
    payload: Mapping[str, Any],
    *,
    algorithm: str = ALG_HMAC,
) -> str:
    """Sign, under the algorithm the caller names rather than one inferred."""

    message = signing_input(domain, payload)
    if algorithm == ALG_HMAC:
        return hmac.new(_key(key), message, hashlib.sha256).hexdigest()
    if algorithm == ALG_ED25519:
        ed25519 = _ed25519_module()
        if not isinstance(key, (bytes, bytearray)) or len(key) != 32:
            raise HandshakeError("an Ed25519 private key is exactly 32 bytes")
        private = ed25519.Ed25519PrivateKey.from_private_bytes(bytes(key))
        return private.sign(message).hex()
    raise HandshakeError(f"unsupported signature algorithm {algorithm!r}")


def verify_payload(
    key: Any,
    domain: str,
    payload: Mapping[str, Any],
    signature: Any,
    *,
    algorithm: str = ALG_HMAC,
) -> bool:
    """Verify under the algorithm the **verifier** expects.
    The algorithm is a parameter and not read from the message.  A verifier that
    trusts the message's own ``alg`` field can be handed a forgery that names
    whichever scheme the attacker can satisfy, which is how the ``alg: none``
    class of bugs works.  Here the caller knows which key it pinned and
    therefore which algorithm is admissible, and a message claiming otherwise
    simply fails.
    """

    if not isinstance(signature, str) or not signature:
        return False
    message = signing_input(domain, payload)
    if algorithm == ALG_HMAC:
        try:
            expected = hmac.new(_key(key), message, hashlib.sha256).hexdigest()
        except HandshakeError:
            return False
        return hmac.compare_digest(expected, signature.lower())
    if algorithm == ALG_ED25519:
        ed25519 = _ed25519_module()
        if not isinstance(key, (bytes, bytearray)) or len(key) != 32:
            return False
        try:
            raw = bytes.fromhex(signature)
        except ValueError:
            return False
        public = ed25519.Ed25519PublicKey.from_public_bytes(bytes(key))
        try:
            public.verify(raw, message)
        except Exception:
            return False
        return True
    raise HandshakeError(f"unsupported signature algorithm {algorithm!r}")


def pair_proof(pairing_code_key: Any, payload: Mapping[str, Any]) -> str:
    """Prove the responder consumed the same one-use pairing code.
    This does **not** authenticate a server at first contact: the agent has just
    sent the code, so whoever received it can compute this.  What it does buy is
    that the pair response cannot be replayed from an earlier session, and that
    the party answering is the party that consumed the code rather than
    something sitting in front of it.  First contact is trust on first use; see
    ``PLAN.md`` section 4.7.
    """

    return sign_payload(pairing_code_key, DOMAIN_PAIR_RESPONSE, payload)


__all__ = [
    "Allowance",
    "DOMAINS",
    "DOMAIN_CATALOGUE",
    "DOMAIN_PAIR_RESPONSE",
    "DOMAIN_RESPONSE",
    "DOMAIN_SNAPSHOT",
    "HANDSHAKE_VERSION",
    "HandshakeError",
    "MAX_COUNTER",
    "MAX_MODELS",
    "MAX_TOOLS",
    "NOTHING",
    "OperatorApproval",
    "RESPONSE_AUTH_KEY",
    "ResponseAuthentication",
    "ALG_ED25519",
    "ALG_HMAC",
    "SIGNATURE_ALGORITHM",
    "SIGNATURE_ALGORITHMS",
    "SNAPSHOT_ALGORITHM",
    "SigningUnavailable",
    "ed25519_public_key",
    "generate_ed25519_key",
    "ServerCatalogue",
    "SignedSnapshot",
    "TOOL_ID_PATTERN",
    "TOOL_VERSION_PATTERN",
    "body_digest",
    "pair_proof",
    "parse_tool_ref",
    "sign_payload",
    "sign_response",
    "signing_input",
    "verify_payload",
]