"""Lab 01 — the agent and workload identity fabric.

The hardest unsolved problem in the JD. Human identity is solved. Service identity is
solved. AGENT identity is neither, because an agent acts on behalf of a user, over
multiple hops, with dynamically discovered tools, sometimes delegating to other agents.

Four properties, each a control an examiner will ask you to demonstrate: credentials are
DERIVED (from a verified assertion, never asserted), NARROWED (at every hop), CHAINED
(the actor list is append-only), and SHORT-LIVED (seconds, never rotated in place).

    pytest test_lab.py -v
    LAB_MODULE=solution pytest test_lab.py -v    # the reference, must be green
"""

from __future__ import annotations

import base64
import hashlib
import hmac
import json
from dataclasses import dataclass, field, replace
from enum import Enum
from typing import Callable, Dict, Iterable, List, Mapping, Optional, Sequence, Set, Tuple

# ======================================================================================
# 1. JOSE: base64url, JWS, and the claim set
# ======================================================================================


def b64url_encode(data: bytes) -> str:
    """Base64url WITHOUT padding, as JOSE requires."""
    # TODO
    raise NotImplementedError


def b64url_decode(text: str) -> bytes:
    """Restore the stripped padding before decoding. Forgetting this produces a binascii
    error on exactly 2/3 of otherwise-valid tokens."""
    # TODO
    raise NotImplementedError


class TokenError(Exception):
    """Verification failed. The message is for the LOG; the caller gets a generic 401 —
    a verifier that says WHICH check failed is an oracle."""

    def __init__(self, reason: str, *, code: str = "invalid_token") -> None:
        super().__init__(reason)
        self.reason = reason
        self.code = code


@dataclass(frozen=True)
class Actor:
    """One link in the delegation chain — RFC 8693's ``act`` claim."""

    sub: str
    kind: str = "agent"          # "agent" | "service" | "workload"


@dataclass(frozen=True)
class Claims:
    """Every field here is checked by somebody. The fields nobody checks are the
    vulnerabilities."""

    iss: str
    sub: str
    aud: str                                  # who it is FOR — the confused-deputy defence
    exp: float
    iat: float
    nbf: float = 0.0
    jti: str = ""
    scope: Tuple[str, ...] = ()
    tenant: str = ""
    act: Tuple[Actor, ...] = ()               # earliest first
    cnf: Optional[str] = None                 # a key thumbprint (RFC 7800)
    client_id: str = ""
    may_delegate: bool = False

    def to_payload(self) -> Dict[str, object]:
        """iss/sub/aud/exp/iat always; the rest ONLY when set, so a minimal token stays
        minimal. ``scope`` is space-delimited (RFC 6749); ``cnf`` is ``{"jkt": ...}``;
        ``act`` is nested by :func:`_nest_actors`."""
        # TODO
        raise NotImplementedError

    @classmethod
    def from_payload(cls, payload: Mapping[str, object]) -> "Claims":
        """The inverse. A missing required claim should raise (the caller converts it)."""
        # TODO
        raise NotImplementedError


def _nest_actors(actors: Sequence[Actor]) -> Dict[str, object]:
    """RFC 8693 NESTS the act claim: the outermost object is the MOST RECENT actor, and
    each ``act`` inside it is the one before. ``actors[0]`` is the earliest.

    ``(a, b)`` -> ``{"sub": "b", "kind": ..., "act": {"sub": "a", "kind": ...}}``
    """
    # TODO
    raise NotImplementedError


def _flatten_actors(value: object) -> Tuple[Actor, ...]:
    """The inverse — earliest first. A missing or malformed value is an empty chain."""
    # TODO
    raise NotImplementedError


class Signer:
    """HMAC-SHA256 JWS. Real deployments use RS256/ES256 so a resource server can verify
    without a minting key; every rule below is identical either way."""

    def __init__(self, key_id: str, secret: bytes, *, algorithm: str = "HS256") -> None:
        self.key_id = key_id
        self.secret = secret
        self.algorithm = algorithm

    def sign(self, claims: Claims) -> str:
        """``b64(header).b64(payload).b64(hmac)`` with header
        ``{"alg": ..., "typ": "at+jwt", "kid": ...}``."""
        # TODO
        raise NotImplementedError

    def verify_signature(self, token: str) -> Tuple[Dict[str, object], Dict[str, object]]:
        """Split into three, decode, then:

        - the header's ``alg`` must equal OURS. Taking the algorithm from the token is
          how "alg: none" and algorithm-confusion attacks work.
        - compare with ``hmac.compare_digest``. A short-circuiting ``==`` leaks the
          signature byte by byte.

        Anything malformed -> TokenError. Returns ``(header, payload)``.
        """
        # TODO
        raise NotImplementedError


def _json_b64(value: Mapping[str, object]) -> str:
    """Canonical JSON — ``sort_keys=True``, no spaces — so identical claims always
    produce an identical token. Not required by JOSE; required by a deterministic test."""
    # TODO
    raise NotImplementedError


# ======================================================================================
# 2. Verification
# ======================================================================================

MAX_CHAIN_DEPTH = 4


@dataclass(frozen=True)
class VerificationPolicy:
    """What THIS resource server requires."""

    audience: str
    trusted_issuers: Tuple[str, ...]
    clock_skew_seconds: float = 30.0
    required_scopes: Tuple[str, ...] = ()
    require_proof_of_possession: bool = False
    max_chain_depth: int = MAX_CHAIN_DEPTH
    max_lifetime_seconds: Optional[float] = None


class ReplayCache:
    """Seen ``jti`` values with expiry. One-time tokens are one-time only if somebody
    remembers."""

    def __init__(self, *, now: Callable[[], float]) -> None:
        self.now = now
        self._seen: Dict[str, float] = {}

    def check_and_record(self, jti: str, expires_at: float) -> None:
        """Evict expired entries first, then raise TokenError(code="replay") if seen."""
        # TODO
        raise NotImplementedError


class Verifier:
    def __init__(self, *, signers: Mapping[str, Signer], policy: VerificationPolicy,
                 now: Callable[[], float],
                 replay_cache: Optional[ReplayCache] = None) -> None:
        self.signers = dict(signers)
        self.policy = policy
        self.now = now
        self.replay_cache = replay_cache

    def verify(self, token: str, *, presented_key: Optional[str] = None) -> Claims:
        """The checks, in order — signature FIRST, because nothing else is trustworthy
        until the token is authentic:

        1. parse the header, select the signer by ``kid`` (unknown -> invalid_key);
        2. verify the signature;
        3. ``iss`` in trusted_issuers                    -> invalid_issuer
        4. ``aud`` == policy.audience                    -> invalid_audience
           **THE confused-deputy defence.** A token minted for another service must not
           be accepted here, or that service can replay it against us.
        5. ``exp``/``nbf``/``iat`` within clock skew     -> expired / not_yet_valid / bad_iat
        6. ``exp - iat`` <= max_lifetime (if set)        -> lifetime_too_long
        7. required scopes present                       -> insufficient_scope
        8. ``len(act)`` <= max_chain_depth               -> chain_too_deep
        9. proof of possession, when required or when ``cnf`` is present:
           no cnf -> pop_required; no key -> pop_missing; mismatch -> pop_mismatch
        10. replay check when ``jti`` is set and a cache is configured
        """
        # TODO
        raise NotImplementedError


def key_thumbprint(public_key: str) -> str:
    """A stand-in for RFC 7638 JWK thumbprinting: sha256, base64url, first 32 chars.
    Deterministic is what matters."""
    # TODO
    raise NotImplementedError


# ======================================================================================
# 3. Scopes — the narrowing algebra
# ======================================================================================


def normalize_scopes(scopes: Iterable[str]) -> Tuple[str, ...]:
    """Sorted, deduplicated, empty strings dropped."""
    # TODO
    raise NotImplementedError


def scope_covers(held: str, requested: str) -> bool:
    """Exact match, or a single trailing wildcard: ``payments.*`` covers
    ``payments.read``. Deliberately limited — a scope language with real pattern matching
    becomes a policy engine, and then nobody can tell what a credential permits by
    reading it."""
    # TODO
    raise NotImplementedError


def narrow_scopes(held: Sequence[str], requested: Sequence[str]) -> Tuple[str, ...]:
    """The intersection, and the ONLY legal direction of travel. This function must be
    incapable of returning a scope not covered by ``held``."""
    # TODO
    raise NotImplementedError


class ScopeEscalation(Exception):
    pass


def require_no_escalation(held: Sequence[str], requested: Sequence[str]) -> Tuple[str, ...]:
    """Narrow, and REFUSE rather than silently dropping.

    Silently returning less than was asked for is worse than an error: the caller
    proceeds believing it has authority it does not, and fails later, far from the cause.
    """
    # TODO
    raise NotImplementedError


# ======================================================================================
# 4. OAuth 2.1
# ======================================================================================


def code_challenge(verifier: str, *, method: str = "S256") -> str:
    """PKCE (RFC 7636): base64url(sha256(verifier)).

    OAuth 2.1 allows **S256 only** — ``plain`` is useless, because an attacker who
    intercepts the challenge has the verifier. Refuse anything else (ValueError), and
    refuse a verifier shorter than 43 characters.
    """
    # TODO
    raise NotImplementedError


@dataclass(frozen=True)
class ClientRegistration:
    client_id: str
    redirect_uris: Tuple[str, ...]
    allowed_scopes: Tuple[str, ...]
    confidential: bool = False
    secret: str = ""


@dataclass
class AuthorizationCode:
    code: str
    client_id: str
    user_id: str
    tenant: str
    scopes: Tuple[str, ...]
    challenge: str
    redirect_uri: str
    expires_at: float
    used: bool = False


class AuthorizationServer:
    """OAuth 2.1 + OIDC, reduced to what an agent platform uses.

    What 2.1 removed, and why: the implicit grant (tokens in a URL fragment leak through
    history and referrers), the password grant (the client sees the password), and bearer
    tokens in query strings. PKCE became mandatory for ALL clients.
    """

    def __init__(self, *, issuer: str, signer: Signer, now: Callable[[], float],
                 clients: Mapping[str, ClientRegistration],
                 code_ttl: float = 60.0, access_ttl: float = 300.0) -> None:
        self.issuer = issuer
        self.signer = signer
        self.now = now
        self.clients = dict(clients)
        self.code_ttl = code_ttl
        self.access_ttl = access_ttl
        self._codes: Dict[str, AuthorizationCode] = {}
        self._counter = 0

    def _next_id(self, prefix: str) -> str:
        self._counter += 1
        return f"{prefix}-{self._counter}"

    def authorize(self, *, client_id: str, redirect_uri: str, scopes: Sequence[str],
                  challenge: str, user_id: str, tenant: str,
                  challenge_method: str = "S256") -> str:
        """Issue an authorization code.

        - unknown client -> invalid_client
        - **EXACT** redirect-URI match -> invalid_redirect. Prefix matching is how an
          open redirect becomes a code-interception attack.
        - method must be S256 -> invalid_challenge_method
        - an empty challenge -> pkce_required
        - scopes narrowed against the client's registration (ScopeEscalation if wider)
        """
        # TODO
        raise NotImplementedError

    def exchange_code(self, *, code: str, client_id: str, verifier: str,
                      redirect_uri: str, audience: str) -> str:
        """Redeem a code for an access token.

        Single use — mark it used BEFORE the expiry check, so a replay is never
        ambiguous. Then: not expired, same client, same redirect_uri, and PKCE
        (``code_challenge(verifier) == stored challenge``). All failures are
        ``invalid_grant``.

        The resulting token carries the user as ``sub``, the tenant, the granted scopes,
        a fresh ``jti``, and ``may_delegate=True``.
        """
        # TODO
        raise NotImplementedError

    def client_credentials(self, *, client_id: str, secret: str,
                           scopes: Sequence[str], audience: str) -> str:
        """Machine-to-machine, no user.

        Simple, and the WRONG default for an agent acting for a person: it erases the
        user from the chain, so every downstream audit record says the platform did it.
        Requires a confidential client and a constant-time secret comparison.
        """
        # TODO
        raise NotImplementedError

    def id_token(self, *, user_id: str, tenant: str, client_id: str) -> str:
        """OIDC's ID token: about the USER, for the CLIENT — so ``aud`` is the client_id,
        not an API. Confusing this with an access token is why people think OIDC and
        OAuth are the same thing."""
        # TODO
        raise NotImplementedError


# ======================================================================================
# 5. RFC 8693 token exchange
# ======================================================================================


class ExchangeError(Exception):
    def __init__(self, reason: str, *, code: str = "invalid_request") -> None:
        super().__init__(reason)
        self.reason = reason
        self.code = code


@dataclass(frozen=True)
class ExchangeRequest:
    subject_token: str
    actor_id: str
    actor_kind: str = "agent"
    audience: str = ""
    scopes: Tuple[str, ...] = ()
    lifetime_seconds: Optional[float] = None
    bind_to_key: Optional[str] = None
    delegation: bool = True           # False = impersonation (chain erased)


class TokenExchange:
    """RFC 8693 — the backbone of multi-hop agent identity."""

    def __init__(self, *, issuer: str, signer: Signer, verifier: Verifier,
                 now: Callable[[], float], max_chain_depth: int = MAX_CHAIN_DEPTH,
                 max_lifetime_seconds: float = 300.0,
                 allow_impersonation: bool = False) -> None:
        self.issuer = issuer
        self.signer = signer
        self.verifier = verifier
        self.now = now
        self.max_chain_depth = max_chain_depth
        self.max_lifetime_seconds = max_lifetime_seconds
        self.allow_impersonation = allow_impersonation
        self._counter = 0

    def exchange(self, request: ExchangeRequest) -> str:
        """VERIFY the subject token (never parse it), then:

        - not ``may_delegate``            -> not_delegable
        - no audience                     -> invalid_request
        - audience unchanged              -> no_narrowing  (an exchange must narrow)
        - impersonation while disabled    -> impersonation_disabled
        - scope not covered               -> scope_escalation
        - the actor already in the chain, or equal to ``sub`` -> chain_cycle
        - resulting chain too deep        -> chain_too_deep
        - no remaining lifetime           -> expired

        The new chain is ``subject.act + (actor,)`` — derived from a VERIFIED assertion,
        never from anything the caller supplied. For impersonation the chain is ERASED
        and ``sub`` becomes the actor, which is why a regulated platform disables it.

        Lifetime is ``min(requested, policy max, the parent's remaining life)`` — a
        derived credential can never outlive its parent. ``may_delegate`` on the result
        is true only while the chain is below the depth limit.
        """
        # TODO
        raise NotImplementedError


# ======================================================================================
# 6. SPIFFE / SPIRE
# ======================================================================================


@dataclass(frozen=True)
class Selector:
    """An attested property of a running workload — the whole point of SPIFFE: identity
    is derived from what the platform can VERIFY, not from a secret."""

    kind: str                 # "k8s:ns" | "k8s:sa" | "docker:image" | "unix:uid"
    value: str


@dataclass(frozen=True)
class SpiffeID:
    trust_domain: str
    path: str

    def __str__(self) -> str:
        return f"spiffe://{self.trust_domain}{self.path}"

    @classmethod
    def parse(cls, value: str) -> "SpiffeID":
        """``spiffe://<trust-domain>/<path>``. Anything else -> ValueError."""
        # TODO
        raise NotImplementedError


@dataclass(frozen=True)
class SVID:
    spiffe_id: SpiffeID
    expires_at: float
    serial: int
    public_key: str

    def is_valid(self, now: float) -> bool:
        # TODO
        raise NotImplementedError


@dataclass(frozen=True)
class RegistrationEntry:
    spiffe_id: str
    selectors: Tuple[Selector, ...]

    def matches(self, attested: Sequence[Selector]) -> bool:
        """ALL registered selectors must be present in the attested set. A subset match
        would let a workload with one matching property claim an identity meant for a
        narrower set."""
        # TODO
        raise NotImplementedError


class SpireServer:
    """Attest a workload, issue an SVID.

    The direction of trust is the lesson: the workload presents NO credential. The
    platform observes verifiable properties and issues an identity from them.
    """

    def __init__(self, *, trust_domain: str, now: Callable[[], float],
                 svid_ttl: float = 60.0) -> None:
        self.trust_domain = trust_domain
        self.now = now
        self.svid_ttl = svid_ttl
        self._entries: List[RegistrationEntry] = []
        self._serial = 0

    def register(self, entry: RegistrationEntry) -> None:
        """Refuse (ValueError) an ID outside this trust domain, or an entry with NO
        selectors — which would match everything."""
        # TODO
        raise NotImplementedError

    def attest(self, attested: Sequence[Selector], *, public_key: str) -> SVID:
        """No match -> attestation_failed. MORE than one match -> ambiguous_attestation:
        guessing would assign an identity nondeterministically, so refuse. Serials
        increase."""
        # TODO
        raise NotImplementedError


@dataclass(frozen=True)
class MtlsPolicy:
    allowed_callers: Tuple[str, ...]
    federated_domains: Tuple[str, ...] = ()


def mtls_authorize(client: SVID, server_policy: MtlsPolicy, *, now: float) -> None:
    """Expired SVID -> svid_expired. Otherwise the caller must be in ``allowed_callers``;
    a caller from another trust domain additionally requires that domain to be federated.
    Federation is EXPLICIT — never a default."""
    # TODO
    raise NotImplementedError


# ======================================================================================
# 7. Non-human identity lifecycle
# ======================================================================================


class NHIState(str, Enum):
    REGISTERED = "registered"
    APPROVED = "approved"
    ACTIVE = "active"
    SUSPENDED = "suspended"
    RETIRED = "retired"


# TODO: the legal transitions.
#   REGISTERED -> APPROVED, RETIRED
#   APPROVED   -> ACTIVE, RETIRED
#   ACTIVE     -> SUSPENDED, RETIRED
#   SUSPENDED  -> ACTIVE, RETIRED
#   RETIRED    -> (terminal)
NHI_TRANSITIONS: Mapping[NHIState, frozenset] = {}


@dataclass(frozen=True)
class NonHumanIdentity:
    """``owner`` is not decoration: an NHI without a named human owner is the finding
    every identity audit produces, because nobody can answer 'should this still
    exist?'"""

    identity_id: str
    kind: str                        # "agent" | "workload" | "service"
    owner: str                       # a HUMAN
    tenant: str
    allowed_scopes: Tuple[str, ...]
    state: NHIState = NHIState.REGISTERED
    max_credential_lifetime: float = 300.0


class IdentityRegistry:
    def __init__(self, *, now: Callable[[], float]) -> None:
        self.now = now
        self._identities: Dict[str, NonHumanIdentity] = {}

    def register(self, identity: NonHumanIdentity) -> NonHumanIdentity:
        """Duplicate id -> ValueError. No owner -> ValueError."""
        # TODO
        raise NotImplementedError

    def get(self, identity_id: str) -> NonHumanIdentity:
        """Unknown -> KeyError."""
        # TODO
        raise NotImplementedError

    def transition(self, identity_id: str, target: NHIState) -> NonHumanIdentity:
        """Undeclared transition -> ValueError."""
        # TODO
        raise NotImplementedError

    def may_receive_credentials(self, identity_id: str) -> bool:
        # TODO
        raise NotImplementedError


# ======================================================================================
# 8. Just-in-time credentials
# ======================================================================================


@dataclass(frozen=True)
class CredentialRequest:
    identity_id: str
    audience: str
    scopes: Tuple[str, ...]
    subject_token: Optional[str] = None      # user context, when acting for someone
    bind_to_key: Optional[str] = None
    lifetime_seconds: float = 60.0


class JitCredentialBroker:
    """Mint at the moment of use, scoped to the action, expiring in seconds.

    Nothing stored, so nothing leaks from storage; short-lived, so a leak is worth
    little; narrowly audienced, so it is worth little *elsewhere*.
    """

    def __init__(self, *, registry: IdentityRegistry, exchange: TokenExchange,
                 issuer: str, signer: Signer, now: Callable[[], float],
                 max_lifetime_seconds: float = 120.0) -> None:
        self.registry = registry
        self.exchange = exchange
        self.issuer = issuer
        self.signer = signer
        self.now = now
        self.max_lifetime_seconds = max_lifetime_seconds
        self._counter = 0

    def issue(self, request: CredentialRequest) -> str:
        """Not ACTIVE -> TokenError(identity_not_active). Revocation latency IS the
        credential TTL, which is why 60 seconds matters.

        Scopes are narrowed against the identity's registered ceiling; the lifetime is
        ``min(requested, broker max, identity max)``.

        With a ``subject_token``: EXCHANGE, so the user stays in the chain.
        Without one: a workload credential with the agent as the sole actor and
        ``may_delegate=False`` — it can do only what the workload may do on its own
        behalf, which is why user-scoped actions must supply a subject token.
        """
        # TODO
        raise NotImplementedError


def describe_chain(claims: Claims) -> str:
    """``sub -> actor1 -> actor2`` — the audit rendering, earliest first."""
    # TODO
    raise NotImplementedError


def main() -> None:
    print("implement the TODOs, then compare with `python solution.py`")


if __name__ == "__main__":
    main()