ietf-draft-analyzer/workspace/packages/act/act/token.py

"""ACT token structures and JWS Compact Serialization.

Defines ACTMandate (Phase 1) and ACTRecord (Phase 2) dataclasses,
plus JWS encoding/decoding primitives for ACT tokens.

Reference: ACT §3 (Lifecycle), §4 (Token Format).
"""

from __future__ import annotations

import base64
import json
import re
import time
import uuid
from dataclasses import dataclass, field
from typing import Any

from .errors import ACTPhaseError, ACTValidationError

# Allowed algorithms per ACT §4.1 — symmetric and "none" are forbidden.
ALLOWED_ALGORITHMS: frozenset[str] = frozenset({"EdDSA", "ES256"})

# Forbidden algorithm prefixes/values per ACT §4.1.
_FORBIDDEN_ALGORITHMS: frozenset[str] = frozenset({
    "none", "HS256", "HS384", "HS512",
})

# Required typ value per ACT §4.1.
ACT_TYP: str = "act+jwt"

# ABNF for action names: component *("." component)
# component = ALPHA *(ALPHA / DIGIT / "-" / "_")
_ACTION_RE = re.compile(
    r"^[A-Za-z][A-Za-z0-9\-_]*(?:\.[A-Za-z][A-Za-z0-9\-_]*)*$"
)


def _b64url_encode(data: bytes) -> str:
    """Base64url encode without padding (RFC 7515 §2)."""
    return base64.urlsafe_b64encode(data).rstrip(b"=").decode("ascii")


def _b64url_decode(s: str) -> bytes:
    """Base64url decode with padding restoration."""
    s = s + "=" * (-len(s) % 4)
    return base64.urlsafe_b64decode(s)


def validate_action_name(action: str) -> None:
    """Validate an action name against ACT ABNF grammar.

    Reference: ACT §4.2.2 (cap action names).

    Raises:
        ACTValidationError: If action does not match the required grammar.
    """
    if not _ACTION_RE.match(action):
        raise ACTValidationError(
            f"Action name {action!r} does not conform to ACT ABNF grammar"
        )


@dataclass(frozen=True)
class TaskClaim:
    """The 'task' claim object.

    Reference: ACT §4.2.2.
    """

    purpose: str
    data_sensitivity: str | None = None
    created_by: str | None = None
    expires_at: int | None = None

    def to_dict(self) -> dict[str, Any]:
        d: dict[str, Any] = {"purpose": self.purpose}
        if self.data_sensitivity is not None:
            d["data_sensitivity"] = self.data_sensitivity
        if self.created_by is not None:
            d["created_by"] = self.created_by
        if self.expires_at is not None:
            d["expires_at"] = self.expires_at
        return d

    @classmethod
    def from_dict(cls, d: dict[str, Any]) -> TaskClaim:
        if "purpose" not in d:
            raise ACTValidationError("task.purpose is required")
        return cls(
            purpose=d["purpose"],
            data_sensitivity=d.get("data_sensitivity"),
            created_by=d.get("created_by"),
            expires_at=d.get("expires_at"),
        )


@dataclass(frozen=True)
class Capability:
    """A single capability entry in the 'cap' array.

    Reference: ACT §4.2.2.
    """

    action: str
    constraints: dict[str, Any] | None = None

    def __post_init__(self) -> None:
        validate_action_name(self.action)

    def to_dict(self) -> dict[str, Any]:
        d: dict[str, Any] = {"action": self.action}
        if self.constraints is not None:
            d["constraints"] = self.constraints
        return d

    @classmethod
    def from_dict(cls, d: dict[str, Any]) -> Capability:
        if "action" not in d:
            raise ACTValidationError("cap[].action is required")
        return cls(
            action=d["action"],
            constraints=d.get("constraints"),
        )


@dataclass(frozen=True)
class DelegationEntry:
    """A single entry in del.chain.

    Reference: ACT §4.2.2 (del), §6 (Delegation Chain).
    """

    delegator: str
    jti: str
    sig: str

    def to_dict(self) -> dict[str, str]:
        return {"delegator": self.delegator, "jti": self.jti, "sig": self.sig}

    @classmethod
    def from_dict(cls, d: dict[str, Any]) -> DelegationEntry:
        for key in ("delegator", "jti", "sig"):
            if key not in d:
                raise ACTValidationError(f"del.chain[].{key} is required")
        return cls(
            delegator=d["delegator"], jti=d["jti"], sig=d["sig"]
        )


@dataclass(frozen=True)
class Delegation:
    """The 'del' claim object.

    Reference: ACT §4.2.2 (del), §6 (Delegation Chain).
    """

    depth: int
    max_depth: int
    chain: list[DelegationEntry] = field(default_factory=list)

    def to_dict(self) -> dict[str, Any]:
        return {
            "depth": self.depth,
            "max_depth": self.max_depth,
            "chain": [e.to_dict() for e in self.chain],
        }

    @classmethod
    def from_dict(cls, d: dict[str, Any]) -> Delegation:
        for key in ("depth", "max_depth"):
            if key not in d:
                raise ACTValidationError(f"del.{key} is required")
        chain_raw = d.get("chain", [])
        chain = [DelegationEntry.from_dict(e) for e in chain_raw]
        return cls(depth=d["depth"], max_depth=d["max_depth"], chain=chain)


@dataclass(frozen=True)
class Oversight:
    """The 'oversight' claim object.

    Reference: ACT §4.2.2 (oversight).
    """

    requires_approval_for: list[str] = field(default_factory=list)
    approval_ref: str | None = None

    def to_dict(self) -> dict[str, Any]:
        d: dict[str, Any] = {
            "requires_approval_for": self.requires_approval_for
        }
        if self.approval_ref is not None:
            d["approval_ref"] = self.approval_ref
        return d

    @classmethod
    def from_dict(cls, d: dict[str, Any]) -> Oversight:
        return cls(
            requires_approval_for=d.get("requires_approval_for", []),
            approval_ref=d.get("approval_ref"),
        )


@dataclass(frozen=True)
class ErrorClaim:
    """The 'err' claim object for failed/partial execution.

    Reference: ACT §4.3.
    """

    code: str
    detail: str

    def to_dict(self) -> dict[str, str]:
        return {"code": self.code, "detail": self.detail}

    @classmethod
    def from_dict(cls, d: dict[str, Any]) -> ErrorClaim:
        for key in ("code", "detail"):
            if key not in d:
                raise ACTValidationError(f"err.{key} is required")
        return cls(code=d["code"], detail=d["detail"])


@dataclass
class ACTMandate:
    """Phase 1 Authorization Mandate.

    Represents a signed authorization from an issuing agent to a target
    agent, encoding capabilities, constraints, and delegation provenance.

    Reference: ACT §3.1, §4.1, §4.2.
    """

    # JOSE header fields
    alg: str
    kid: str
    x5c: list[str] | None = None
    did: str | None = None

    # Required JWT claims
    iss: str = ""
    sub: str = ""
    aud: str | list[str] = ""
    iat: int = 0
    exp: int = 0
    jti: str = field(default_factory=lambda: str(uuid.uuid4()))

    # Optional standard claims
    wid: str | None = None

    # Required ACT claims
    task: TaskClaim = field(default_factory=lambda: TaskClaim(purpose=""))
    cap: list[Capability] = field(default_factory=list)

    # Optional ACT claims
    delegation: Delegation | None = None
    oversight: Oversight | None = None

    def validate(self) -> None:
        """Validate structural correctness of this mandate.

        Reference: ACT §4.1, §4.2, §8.1 step 11.

        Raises:
            ACTValidationError: If any required field is missing or invalid.
        """
        _validate_algorithm(self.alg)
        if not self.kid:
            raise ACTValidationError("kid is required in JOSE header")
        for claim_name in ("iss", "sub", "aud", "jti"):
            val = getattr(self, claim_name)
            if not val:
                raise ACTValidationError(f"{claim_name} claim is required")
        if self.iat <= 0:
            raise ACTValidationError("iat must be a positive NumericDate")
        if self.exp <= 0:
            raise ACTValidationError("exp must be a positive NumericDate")
        if not self.task.purpose:
            raise ACTValidationError("task.purpose is required")
        if not self.cap:
            raise ACTValidationError("cap must contain at least one capability")

    def to_header(self) -> dict[str, Any]:
        """Build JOSE header dict.

        Reference: ACT §4.1.
        """
        h: dict[str, Any] = {
            "alg": self.alg,
            "typ": ACT_TYP,
            "kid": self.kid,
        }
        if self.x5c is not None:
            h["x5c"] = self.x5c
        if self.did is not None:
            h["did"] = self.did
        return h

    def to_claims(self) -> dict[str, Any]:
        """Build JWT claims dict (Phase 1 claims only).

        Reference: ACT §4.2.
        """
        c: dict[str, Any] = {
            "iss": self.iss,
            "sub": self.sub,
            "aud": self.aud,
            "iat": self.iat,
            "exp": self.exp,
            "jti": self.jti,
            "task": self.task.to_dict(),
            "cap": [cap.to_dict() for cap in self.cap],
        }
        if self.wid is not None:
            c["wid"] = self.wid
        if self.delegation is not None:
            c["del"] = self.delegation.to_dict()
        if self.oversight is not None:
            c["oversight"] = self.oversight.to_dict()
        return c

    def signing_input(self) -> bytes:
        """Compute the JWS signing input (header.payload) as bytes.

        Reference: RFC 7515 §5.1.
        """
        header_b64 = _b64url_encode(
            json.dumps(self.to_header(), separators=(",", ":")).encode()
        )
        payload_b64 = _b64url_encode(
            json.dumps(self.to_claims(), separators=(",", ":")).encode()
        )
        return f"{header_b64}.{payload_b64}".encode("ascii")

    def is_phase2(self) -> bool:
        """Return False; mandates are always Phase 1."""
        return False

    @classmethod
    def from_claims(
        cls,
        header: dict[str, Any],
        claims: dict[str, Any],
    ) -> ACTMandate:
        """Construct an ACTMandate from parsed header and claims dicts.

        Reference: ACT §4.1, §4.2.

        Raises:
            ACTValidationError: If required fields are missing.
            ACTPhaseError: If exec_act is present (this is a Phase 2 token).
        """
        if "exec_act" in claims:
            raise ACTPhaseError(
                "Token contains exec_act; use ACTRecord.from_claims instead"
            )

        del_raw = claims.get("del")
        delegation = Delegation.from_dict(del_raw) if del_raw else None

        oversight_raw = claims.get("oversight")
        oversight_obj = Oversight.from_dict(oversight_raw) if oversight_raw else None

        task_raw = claims.get("task")
        if task_raw is None:
            raise ACTValidationError("task claim is required")

        cap_raw = claims.get("cap")
        if cap_raw is None:
            raise ACTValidationError("cap claim is required")

        return cls(
            alg=header.get("alg", ""),
            kid=header.get("kid", ""),
            x5c=header.get("x5c"),
            did=header.get("did"),
            iss=claims.get("iss", ""),
            sub=claims.get("sub", ""),
            aud=claims.get("aud", ""),
            iat=claims.get("iat", 0),
            exp=claims.get("exp", 0),
            jti=claims.get("jti", ""),
            wid=claims.get("wid"),
            task=TaskClaim.from_dict(task_raw),
            cap=[Capability.from_dict(c) for c in cap_raw],
            delegation=delegation,
            oversight=oversight_obj,
        )


@dataclass
class ACTRecord:
    """Phase 2 Execution Record.

    Contains all Phase 1 claims preserved unchanged, plus execution
    claims added by the executing agent. Re-signed by sub's key.

    Reference: ACT §3.2, §4.3.
    """

    # JOSE header fields (Phase 2 header uses sub's kid)
    alg: str
    kid: str
    x5c: list[str] | None = None
    did: str | None = None

    # Phase 1 claims (preserved)
    iss: str = ""
    sub: str = ""
    aud: str | list[str] = ""
    iat: int = 0
    exp: int = 0
    jti: str = ""
    wid: str | None = None
    task: TaskClaim = field(default_factory=lambda: TaskClaim(purpose=""))
    cap: list[Capability] = field(default_factory=list)
    delegation: Delegation | None = None
    oversight: Oversight | None = None

    # Phase 2 claims (execution)
    exec_act: str = ""
    pred: list[str] = field(default_factory=list)
    exec_ts: int = 0
    status: str = ""
    inp_hash: str | None = None
    out_hash: str | None = None
    err: ErrorClaim | None = None

    def validate(self) -> None:
        """Validate structural correctness of this record.

        Reference: ACT §4.3, §8.2 steps 13-16.

        Raises:
            ACTValidationError: If any required field is missing or invalid.
        """
        _validate_algorithm(self.alg)
        if not self.kid:
            raise ACTValidationError("kid is required in JOSE header")
        for claim_name in ("iss", "sub", "aud", "jti"):
            val = getattr(self, claim_name)
            if not val:
                raise ACTValidationError(f"{claim_name} claim is required")
        if self.iat <= 0:
            raise ACTValidationError("iat must be a positive NumericDate")
        if self.exp <= 0:
            raise ACTValidationError("exp must be a positive NumericDate")
        if not self.task.purpose:
            raise ACTValidationError("task.purpose is required")
        if not self.cap:
            raise ACTValidationError("cap must contain at least one capability")
        if not self.exec_act:
            raise ACTValidationError("exec_act is required in Phase 2")
        validate_action_name(self.exec_act)
        if self.exec_ts <= 0:
            raise ACTValidationError("exec_ts must be a positive NumericDate")
        if self.status not in ("completed", "failed", "partial"):
            raise ACTValidationError(
                f"status must be one of completed/failed/partial, got {self.status!r}"
            )

    def to_header(self) -> dict[str, Any]:
        """Build JOSE header dict for Phase 2.

        In Phase 2, kid MUST reference the sub agent's key.
        Reference: ACT §4.1, §8.2 step 17.
        """
        h: dict[str, Any] = {
            "alg": self.alg,
            "typ": ACT_TYP,
            "kid": self.kid,
        }
        if self.x5c is not None:
            h["x5c"] = self.x5c
        if self.did is not None:
            h["did"] = self.did
        return h

    def to_claims(self) -> dict[str, Any]:
        """Build JWT claims dict (Phase 1 + Phase 2 claims).

        Reference: ACT §4.2, §4.3.
        """
        c: dict[str, Any] = {
            "iss": self.iss,
            "sub": self.sub,
            "aud": self.aud,
            "iat": self.iat,
            "exp": self.exp,
            "jti": self.jti,
            "task": self.task.to_dict(),
            "cap": [cap.to_dict() for cap in self.cap],
            "exec_act": self.exec_act,
            "pred": self.pred,
            "exec_ts": self.exec_ts,
            "status": self.status,
        }
        if self.wid is not None:
            c["wid"] = self.wid
        if self.delegation is not None:
            c["del"] = self.delegation.to_dict()
        if self.oversight is not None:
            c["oversight"] = self.oversight.to_dict()
        if self.inp_hash is not None:
            c["inp_hash"] = self.inp_hash
        if self.out_hash is not None:
            c["out_hash"] = self.out_hash
        if self.err is not None:
            c["err"] = self.err.to_dict()
        return c

    def signing_input(self) -> bytes:
        """Compute the JWS signing input (header.payload) as bytes.

        Reference: RFC 7515 §5.1.
        """
        header_b64 = _b64url_encode(
            json.dumps(self.to_header(), separators=(",", ":")).encode()
        )
        payload_b64 = _b64url_encode(
            json.dumps(self.to_claims(), separators=(",", ":")).encode()
        )
        return f"{header_b64}.{payload_b64}".encode("ascii")

    def is_phase2(self) -> bool:
        """Return True; records are always Phase 2."""
        return True

    @classmethod
    def from_mandate(
        cls,
        mandate: ACTMandate,
        *,
        kid: str,
        exec_act: str,
        pred: list[str] | None = None,
        exec_ts: int | None = None,
        status: str = "completed",
        inp_hash: str | None = None,
        out_hash: str | None = None,
        err: ErrorClaim | None = None,
    ) -> ACTRecord:
        """Create an ACTRecord by transitioning a mandate to Phase 2.

        The kid MUST be the sub agent's key identifier.

        Reference: ACT §3.2, §4.3.
        """
        return cls(
            alg=mandate.alg,
            kid=kid,
            x5c=mandate.x5c,
            did=mandate.did,
            iss=mandate.iss,
            sub=mandate.sub,
            aud=mandate.aud,
            iat=mandate.iat,
            exp=mandate.exp,
            jti=mandate.jti,
            wid=mandate.wid,
            task=mandate.task,
            cap=mandate.cap,
            delegation=mandate.delegation,
            oversight=mandate.oversight,
            exec_act=exec_act,
            pred=pred if pred is not None else [],
            exec_ts=exec_ts if exec_ts is not None else int(time.time()),
            status=status,
            inp_hash=inp_hash,
            out_hash=out_hash,
            err=err,
        )

    @classmethod
    def from_claims(
        cls,
        header: dict[str, Any],
        claims: dict[str, Any],
    ) -> ACTRecord:
        """Construct an ACTRecord from parsed header and claims dicts.

        Reference: ACT §4.1, §4.2, §4.3.

        Raises:
            ACTValidationError: If required fields are missing.
            ACTPhaseError: If exec_act is absent (this is a Phase 1 token).
        """
        if "exec_act" not in claims:
            raise ACTPhaseError(
                "Token does not contain exec_act; use ACTMandate.from_claims instead"
            )

        del_raw = claims.get("del")
        delegation = Delegation.from_dict(del_raw) if del_raw else None

        oversight_raw = claims.get("oversight")
        oversight_obj = Oversight.from_dict(oversight_raw) if oversight_raw else None

        task_raw = claims.get("task")
        if task_raw is None:
            raise ACTValidationError("task claim is required")

        cap_raw = claims.get("cap")
        if cap_raw is None:
            raise ACTValidationError("cap claim is required")

        err_raw = claims.get("err")
        err_obj = ErrorClaim.from_dict(err_raw) if err_raw else None

        return cls(
            alg=header.get("alg", ""),
            kid=header.get("kid", ""),
            x5c=header.get("x5c"),
            did=header.get("did"),
            iss=claims.get("iss", ""),
            sub=claims.get("sub", ""),
            aud=claims.get("aud", ""),
            iat=claims.get("iat", 0),
            exp=claims.get("exp", 0),
            jti=claims.get("jti", ""),
            wid=claims.get("wid"),
            task=TaskClaim.from_dict(task_raw),
            cap=[Capability.from_dict(c) for c in cap_raw],
            delegation=delegation,
            oversight=oversight_obj,
            exec_act=claims["exec_act"],
            pred=claims.get("pred", []),
            exec_ts=claims.get("exec_ts", 0),
            status=claims.get("status", ""),
            inp_hash=claims.get("inp_hash"),
            out_hash=claims.get("out_hash"),
            err=err_obj,
        )


# --- JWS Compact Serialization ---


def encode_jws(
    token: ACTMandate | ACTRecord,
    signature: bytes,
) -> str:
    """Encode a token and signature as JWS Compact Serialization.

    Returns header.payload.signature (three base64url segments).

    Reference: RFC 7515 §3.1, ACT §4.
    """
    signing_input = token.signing_input().decode("ascii")
    sig_b64 = _b64url_encode(signature)
    return f"{signing_input}.{sig_b64}"


def decode_jws(compact: str) -> tuple[dict[str, Any], dict[str, Any], bytes, bytes]:
    """Decode a JWS Compact Serialization string.

    Returns (header_dict, claims_dict, signature_bytes, signing_input_bytes).

    Reference: RFC 7515 §5.2, ACT §4.

    Raises:
        ACTValidationError: If the token is malformed.
    """
    parts = compact.split(".")
    if len(parts) != 3:
        raise ACTValidationError(
            f"JWS Compact Serialization requires 3 parts, got {len(parts)}"
        )

    try:
        header = json.loads(_b64url_decode(parts[0]))
    except (json.JSONDecodeError, Exception) as e:
        raise ACTValidationError(f"Invalid JOSE header: {e}") from e

    try:
        claims = json.loads(_b64url_decode(parts[1]))
    except (json.JSONDecodeError, Exception) as e:
        raise ACTValidationError(f"Invalid JWT claims: {e}") from e

    try:
        signature = _b64url_decode(parts[2])
    except Exception as e:
        raise ACTValidationError(f"Invalid signature encoding: {e}") from e

    signing_input = f"{parts[0]}.{parts[1]}".encode("ascii")

    # Validate header requirements per ACT §4.1
    typ = header.get("typ")
    if typ != ACT_TYP:
        raise ACTValidationError(
            f"typ must be {ACT_TYP!r}, got {typ!r}"
        )

    alg = header.get("alg", "")
    _validate_algorithm(alg)

    if "kid" not in header:
        raise ACTValidationError("kid is required in JOSE header")

    return header, claims, signature, signing_input


def parse_token(compact: str) -> ACTMandate | ACTRecord:
    """Parse a JWS compact string into an ACTMandate or ACTRecord.

    Determines phase by presence of exec_act claim.

    Reference: ACT §3 (phase determination).

    Returns:
        ACTMandate for Phase 1, ACTRecord for Phase 2.
    """
    header, claims, _, _ = decode_jws(compact)
    if "exec_act" in claims:
        return ACTRecord.from_claims(header, claims)
    return ACTMandate.from_claims(header, claims)


def _validate_algorithm(alg: str) -> None:
    """Check algorithm is allowed per ACT §4.1.

    Raises:
        ACTValidationError: If algorithm is forbidden or unsupported.
    """
    if alg in _FORBIDDEN_ALGORITHMS or alg.upper() in _FORBIDDEN_ALGORITHMS:
        raise ACTValidationError(
            f"Algorithm {alg!r} is forbidden by ACT specification"
        )
    if alg not in ALLOWED_ALGORITHMS:
        raise ACTValidationError(
            f"Unsupported algorithm {alg!r}; allowed: {sorted(ALLOWED_ALGORITHMS)}"
        )