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

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

import base64
import importlib
import json
import os

import pytest

lab = importlib.import_module(os.environ.get("LAB_MODULE", "lab"))


def clock(start=1000.0, step=1.0):
    state = {"t": start - step}

    def now():
        state["t"] += step
        return state["t"]

    return now


def frozen(value=1000.0):
    return lambda: value


def payload_of(token):
    return json.loads(lab.b64url_decode(token.split(".")[1]))


def claims_of(token):
    return lab.Claims.from_payload(payload_of(token))


def signer(key_id="k1", secret=b"secret"):
    return lab.Signer(key_id, secret)


def claims(**kwargs):
    kwargs.setdefault("iss", "https://login.bank.ae")
    kwargs.setdefault("sub", "u-42")
    kwargs.setdefault("aud", "agent-platform")
    kwargs.setdefault("iat", 1000.0)
    kwargs.setdefault("exp", 1300.0)
    return lab.Claims(**kwargs)


def policy(**kwargs):
    kwargs.setdefault("audience", "agent-platform")
    kwargs.setdefault("trusted_issuers", ("https://login.bank.ae",))
    return lab.VerificationPolicy(**kwargs)


# ======================================================================================
# 1. Encoding and signing
# ======================================================================================


@pytest.mark.parametrize("data", [b"", b"a", b"ab", b"abc", b"abcd", b"\xff\xfe\x00"])
def test_base64url_round_trips_every_padding_case(data):
    assert lab.b64url_decode(lab.b64url_encode(data)) == data


def test_base64url_output_is_unpadded():
    assert "=" not in lab.b64url_encode(b"abcde")


def test_a_token_has_three_parts():
    assert len(signer().sign(claims()).split(".")) == 3


def test_signing_is_deterministic():
    """Canonical JSON: the same claims always produce the same token, which is what
    makes every downstream test an equality assertion."""
    assert signer().sign(claims()) == signer().sign(claims())


def test_a_tampered_payload_fails_verification():
    token = signer().sign(claims())
    header, _, signature = token.split(".")
    forged = lab.b64url_encode(json.dumps({**payload_of(token), "sub": "attacker"}).encode())
    with pytest.raises(lab.TokenError):
        signer().verify_signature(f"{header}.{forged}.{signature}")


def test_a_different_key_fails_verification():
    token = lab.Signer("k1", b"real").sign(claims())
    with pytest.raises(lab.TokenError):
        lab.Signer("k1", b"attacker").verify_signature(token)


def test_an_unexpected_algorithm_is_refused():
    """'alg: none' and algorithm confusion are the two classic JWT breaks."""
    s = signer()
    header = lab.b64url_encode(json.dumps({"alg": "none", "kid": "k1"}).encode())
    body = lab.b64url_encode(json.dumps(claims().to_payload()).encode())
    with pytest.raises(lab.TokenError) as exc:
        s.verify_signature(f"{header}.{body}.")
    assert exc.value.code == "invalid_alg"


@pytest.mark.parametrize("token", ["", "a", "a.b", "a.b.c.d", "!!!.???.$$$"])
def test_malformed_tokens_are_refused(token):
    with pytest.raises(lab.TokenError):
        signer().verify_signature(token)


# ======================================================================================
# 2. Claims and the actor chain
# ======================================================================================


def test_claims_round_trip():
    original = claims(scope=("a", "b"), tenant="wholesale", jti="j1", nbf=999.0,
                      client_id="c1", may_delegate=True, cnf="thumb")
    assert lab.Claims.from_payload(original.to_payload()) == original


def test_scope_is_space_delimited_on_the_wire():
    assert claims(scope=("payments.read", "crm.read")).to_payload()["scope"] == \
        "payments.read crm.read"


def test_optional_claims_are_omitted_when_empty():
    payload = claims().to_payload()
    for key in ("nbf", "jti", "scope", "tenant", "cnf", "act", "may_delegate"):
        assert key not in payload


def test_the_actor_chain_nests_with_the_latest_outermost():
    """RFC 8693 nests `act`: outermost is the MOST RECENT actor — the opposite of what
    most people assume."""
    payload = claims(act=(lab.Actor("orchestrator"), lab.Actor("investigator"))).to_payload()
    assert payload["act"]["sub"] == "investigator"
    assert payload["act"]["act"]["sub"] == "orchestrator"


def test_the_actor_chain_round_trips_earliest_first():
    chain = (lab.Actor("a"), lab.Actor("b"), lab.Actor("c"))
    assert lab.Claims.from_payload(claims(act=chain).to_payload()).act == chain


def test_an_absent_chain_is_empty():
    assert lab.Claims.from_payload(claims().to_payload()).act == ()


def test_describe_chain_renders_user_then_actors():
    assert lab.describe_chain(claims(act=(lab.Actor("o"), lab.Actor("i")))) == \
        "u-42 -> o -> i"


# ======================================================================================
# 3. Verification
# ======================================================================================


def verifier(signers=None, pol=None, now=None, cache=None):
    return lab.Verifier(signers=signers or {"k1": signer()}, policy=pol or policy(),
                        now=now or frozen(1100.0), replay_cache=cache)


def test_a_good_token_verifies():
    got = verifier().verify(signer().sign(claims(scope=("a",))))
    assert got.sub == "u-42"
    assert got.scope == ("a",)


def test_an_untrusted_issuer_is_refused():
    with pytest.raises(lab.TokenError) as exc:
        verifier().verify(signer().sign(claims(iss="https://evil.example")))
    assert exc.value.code == "invalid_issuer"


def test_a_token_for_another_audience_is_refused():
    """The confused-deputy defence: the single most important check."""
    with pytest.raises(lab.TokenError) as exc:
        verifier().verify(signer().sign(claims(aud="core-banking")))
    assert exc.value.code == "invalid_audience"


def test_an_expired_token_is_refused():
    with pytest.raises(lab.TokenError) as exc:
        verifier(now=frozen(2000.0)).verify(signer().sign(claims()))
    assert exc.value.code == "expired"


def test_clock_skew_is_tolerated_in_both_directions():
    v = verifier(now=frozen(1320.0), pol=policy(clock_skew_seconds=30.0))
    v.verify(signer().sign(claims(exp=1300.0)))          # 20s past expiry, within skew
    early = verifier(now=frozen(980.0), pol=policy(clock_skew_seconds=30.0))
    early.verify(signer().sign(claims(nbf=1000.0)))      # 20s before nbf, within skew


def test_skew_has_a_limit():
    with pytest.raises(lab.TokenError):
        verifier(now=frozen(1400.0), pol=policy(clock_skew_seconds=30.0)).verify(
            signer().sign(claims(exp=1300.0)))


def test_a_not_yet_valid_token_is_refused():
    with pytest.raises(lab.TokenError) as exc:
        verifier(now=frozen(1000.0)).verify(signer().sign(claims(nbf=1200.0, exp=1500.0)))
    assert exc.value.code == "not_yet_valid"


def test_an_unknown_key_id_is_refused():
    with pytest.raises(lab.TokenError) as exc:
        verifier().verify(lab.Signer("other-kid", b"secret").sign(claims()))
    assert exc.value.code == "invalid_key"


def test_missing_required_scopes_are_refused():
    with pytest.raises(lab.TokenError) as exc:
        verifier(pol=policy(required_scopes=("payments.release",))).verify(
            signer().sign(claims(scope=("payments.read",))))
    assert exc.value.code == "insufficient_scope"


def test_an_over_long_lifetime_is_refused():
    """A short-lived credential minted with a long lifetime is not short-lived."""
    with pytest.raises(lab.TokenError) as exc:
        verifier(pol=policy(max_lifetime_seconds=60.0)).verify(
            signer().sign(claims(iat=1000.0, exp=1300.0)))
    assert exc.value.code == "lifetime_too_long"


def test_an_over_deep_chain_is_refused():
    chain = tuple(lab.Actor(f"a{i}") for i in range(5))
    with pytest.raises(lab.TokenError) as exc:
        verifier(pol=policy(max_chain_depth=4)).verify(signer().sign(claims(act=chain)))
    assert exc.value.code == "chain_too_deep"


def test_proof_of_possession_binds_the_token_to_a_key():
    token = signer().sign(claims(cnf=lab.key_thumbprint("pk-good")))
    v = verifier(pol=policy(require_proof_of_possession=True))
    v.verify(token, presented_key="pk-good")
    with pytest.raises(lab.TokenError) as exc:
        v.verify(token, presented_key="pk-stolen")
    assert exc.value.code == "pop_mismatch"


def test_a_bound_token_without_a_key_is_refused():
    token = signer().sign(claims(cnf=lab.key_thumbprint("pk")))
    with pytest.raises(lab.TokenError) as exc:
        verifier().verify(token)
    assert exc.value.code == "pop_missing"


def test_a_policy_requiring_pop_refuses_an_unbound_token():
    with pytest.raises(lab.TokenError) as exc:
        verifier(pol=policy(require_proof_of_possession=True)).verify(
            signer().sign(claims()))
    assert exc.value.code == "pop_required"


def test_key_thumbprints_are_deterministic_and_distinct():
    assert lab.key_thumbprint("a") == lab.key_thumbprint("a")
    assert lab.key_thumbprint("a") != lab.key_thumbprint("b")


def test_replay_is_detected_and_expired_entries_are_evicted():
    cache = lab.ReplayCache(now=frozen(1100.0))
    v = verifier(cache=cache)
    token = signer().sign(claims(jti="j1"))
    v.verify(token)
    with pytest.raises(lab.TokenError) as exc:
        v.verify(token)
    assert exc.value.code == "replay"

    later = lab.ReplayCache(now=frozen(2000.0))
    later.check_and_record("j1", 1500.0)
    later.check_and_record("j1", 2500.0)          # the first entry had expired


def test_a_token_without_a_jti_is_not_replay_checked():
    v = verifier(cache=lab.ReplayCache(now=frozen(1100.0)))
    token = signer().sign(claims())
    v.verify(token)
    v.verify(token)                                # no jti, no replay protection


# ======================================================================================
# 4. Scopes
# ======================================================================================


def test_exact_and_wildcard_coverage():
    assert lab.scope_covers("payments.read", "payments.read")
    assert lab.scope_covers("payments.*", "payments.read")
    assert not lab.scope_covers("payments.read", "payments.release")
    assert not lab.scope_covers("payments.*", "crm.read")


def test_narrowing_is_an_intersection():
    assert lab.narrow_scopes(["a", "b"], ["b", "c"]) == ("b",)


def test_narrowing_can_only_shrink():
    """The property that makes escalation structurally impossible."""
    held = ["payments.read"]
    for requested in (["payments.release"], ["treasury.trade"], ["*"]):
        assert set(lab.narrow_scopes(held, requested)) <= set(held) | set()


def test_narrowing_is_normalized_and_deduplicated():
    assert lab.narrow_scopes(["a"], ["a", "a", ""]) == ("a",)


def test_require_no_escalation_refuses_rather_than_silently_dropping():
    """Silently granting less is worse than an error: the caller proceeds believing it
    has authority it does not."""
    with pytest.raises(lab.ScopeEscalation):
        lab.require_no_escalation(["payments.read"], ["payments.read", "payments.release"])
    assert lab.require_no_escalation(["payments.*"], ["payments.read"]) == ("payments.read",)


# ======================================================================================
# 5. OAuth 2.1
# ======================================================================================


def clients():
    return {
        "web": lab.ClientRegistration("web", ("https://bank.ae/cb",),
                                      ("payments.read", "payments.release")),
        "batch": lab.ClientRegistration("batch", (), ("reports.generate",),
                                        confidential=True, secret="s3cret"),
    }


def auth_server(now=None):
    return lab.AuthorizationServer(issuer="https://login.bank.ae", signer=signer(),
                                   now=now or clock(), clients=clients())


VERIFIER_SECRET = "v" * 43


def test_pkce_challenge_is_s256_and_deterministic():
    assert lab.code_challenge(VERIFIER_SECRET) == lab.code_challenge(VERIFIER_SECRET)
    assert lab.code_challenge(VERIFIER_SECRET) != lab.code_challenge("w" * 43)


def test_plain_pkce_is_refused():
    """OAuth 2.1 allows S256 only: with `plain`, intercepting the challenge gives you the
    verifier."""
    with pytest.raises(ValueError):
        lab.code_challenge(VERIFIER_SECRET, method="plain")


def test_a_short_verifier_is_refused():
    with pytest.raises(ValueError):
        lab.code_challenge("short")


def test_the_authorization_code_flow():
    auth = auth_server()
    code = auth.authorize(client_id="web", redirect_uri="https://bank.ae/cb",
                          scopes=["payments.read"],
                          challenge=lab.code_challenge(VERIFIER_SECRET),
                          user_id="u-42", tenant="wholesale")
    token = auth.exchange_code(code=code, client_id="web", verifier=VERIFIER_SECRET,
                               redirect_uri="https://bank.ae/cb", audience="agent-platform")
    got = claims_of(token)
    assert got.sub == "u-42"
    assert got.aud == "agent-platform"
    assert got.tenant == "wholesale"
    assert got.may_delegate is True


def test_pkce_verification_failure():
    auth = auth_server()
    code = auth.authorize(client_id="web", redirect_uri="https://bank.ae/cb",
                          scopes=["payments.read"],
                          challenge=lab.code_challenge(VERIFIER_SECRET),
                          user_id="u-42", tenant="wholesale")
    with pytest.raises(lab.TokenError) as exc:
        auth.exchange_code(code=code, client_id="web", verifier="x" * 43,
                           redirect_uri="https://bank.ae/cb", audience="agent-platform")
    assert exc.value.code == "invalid_grant"


def test_an_authorization_code_is_single_use():
    auth = auth_server()
    code = auth.authorize(client_id="web", redirect_uri="https://bank.ae/cb",
                          scopes=["payments.read"],
                          challenge=lab.code_challenge(VERIFIER_SECRET),
                          user_id="u-42", tenant="wholesale")
    auth.exchange_code(code=code, client_id="web", verifier=VERIFIER_SECRET,
                       redirect_uri="https://bank.ae/cb", audience="agent-platform")
    with pytest.raises(lab.TokenError):
        auth.exchange_code(code=code, client_id="web", verifier=VERIFIER_SECRET,
                           redirect_uri="https://bank.ae/cb", audience="agent-platform")


def test_redirect_uri_matching_is_exact():
    auth = auth_server()
    with pytest.raises(lab.TokenError) as exc:
        auth.authorize(client_id="web", redirect_uri="https://bank.ae/cb/evil",
                       scopes=["payments.read"],
                       challenge=lab.code_challenge(VERIFIER_SECRET),
                       user_id="u-42", tenant="wholesale")
    assert exc.value.code == "invalid_redirect"


def test_pkce_is_mandatory():
    auth = auth_server()
    with pytest.raises(lab.TokenError) as exc:
        auth.authorize(client_id="web", redirect_uri="https://bank.ae/cb",
                       scopes=["payments.read"], challenge="",
                       user_id="u-42", tenant="wholesale")
    assert exc.value.code == "pkce_required"


def test_a_code_cannot_be_redeemed_by_another_client():
    auth = auth_server()
    code = auth.authorize(client_id="web", redirect_uri="https://bank.ae/cb",
                          scopes=["payments.read"],
                          challenge=lab.code_challenge(VERIFIER_SECRET),
                          user_id="u-42", tenant="wholesale")
    with pytest.raises(lab.TokenError):
        auth.exchange_code(code=code, client_id="batch", verifier=VERIFIER_SECRET,
                           redirect_uri="https://bank.ae/cb", audience="agent-platform")


def test_a_client_cannot_request_scopes_it_is_not_registered_for():
    with pytest.raises(lab.ScopeEscalation):
        auth_server().authorize(client_id="web", redirect_uri="https://bank.ae/cb",
                                scopes=["treasury.trade"],
                                challenge=lab.code_challenge(VERIFIER_SECRET),
                                user_id="u-42", tenant="wholesale")


def test_client_credentials_requires_a_confidential_client_and_the_secret():
    auth = auth_server()
    token = auth.client_credentials(client_id="batch", secret="s3cret",
                                    scopes=["reports.generate"], audience="reporting")
    assert claims_of(token).sub == "batch"
    with pytest.raises(lab.TokenError):
        auth.client_credentials(client_id="batch", secret="wrong",
                                scopes=["reports.generate"], audience="reporting")
    with pytest.raises(lab.TokenError):
        auth.client_credentials(client_id="web", secret="",
                                scopes=["payments.read"], audience="reporting")


def test_a_client_credentials_token_has_no_user_and_no_chain():
    """Which is exactly why it is the wrong default for an agent acting for a person."""
    token = auth_server().client_credentials(
        client_id="batch", secret="s3cret", scopes=["reports.generate"],
        audience="reporting")
    got = claims_of(token)
    assert got.act == ()
    assert got.may_delegate is False


def test_an_id_token_is_audienced_to_the_client_not_the_api():
    token = auth_server().id_token(user_id="u-42", tenant="wholesale", client_id="web")
    assert claims_of(token).aud == "web"


# ======================================================================================
# 6. Token exchange
# ======================================================================================


def exchange_setup(*, audience="agent-platform", max_depth=4, max_lifetime=120.0,
                   allow_impersonation=False, now=None):
    now = now or frozen(1100.0)
    s = signer()
    v = lab.Verifier(signers={"k1": s},
                     policy=lab.VerificationPolicy(
                         audience=audience,
                         trusted_issuers=("https://login.bank.ae", "https://sts.bank.ae")),
                     now=now)
    return lab.TokenExchange(issuer="https://sts.bank.ae", signer=s, verifier=v, now=now,
                             max_chain_depth=max_depth, max_lifetime_seconds=max_lifetime,
                             allow_impersonation=allow_impersonation), s


def user_token(s, **kwargs):
    kwargs.setdefault("scope", ("payments.read", "payments.release"))
    kwargs.setdefault("may_delegate", True)
    kwargs.setdefault("tenant", "wholesale")
    kwargs.setdefault("exp", 2000.0)
    return s.sign(claims(**kwargs))


def test_an_exchange_narrows_the_audience_and_appends_to_the_chain():
    ex, s = exchange_setup()
    out = ex.exchange(lab.ExchangeRequest(
        subject_token=user_token(s), actor_id="orchestrator",
        audience="investigator", scopes=("payments.read",)))
    got = claims_of(out)
    assert got.aud == "investigator"
    assert got.sub == "u-42"                       # the USER remains the subject
    assert [a.sub for a in got.act] == ["orchestrator"]
    assert got.scope == ("payments.read",)


def test_the_chain_is_derived_from_a_verified_token_not_from_the_request():
    """A callee cannot assert its own position; the chain comes from the subject token."""
    ex, s = exchange_setup()
    first = ex.exchange(lab.ExchangeRequest(
        subject_token=user_token(s), actor_id="orchestrator",
        audience="investigator", scopes=("payments.read",)))
    ex2, _ = exchange_setup(audience="investigator")
    second = ex2.exchange(lab.ExchangeRequest(
        subject_token=first, actor_id="investigator",
        audience="core-banking", scopes=("payments.read",)))
    assert [a.sub for a in claims_of(second).act] == ["orchestrator", "investigator"]


def test_an_exchange_cannot_widen_scope():
    ex, s = exchange_setup()
    with pytest.raises(lab.ExchangeError) as exc:
        ex.exchange(lab.ExchangeRequest(
            subject_token=user_token(s, scope=("payments.read",)),
            actor_id="orchestrator", audience="investigator",
            scopes=("payments.release",)))
    assert exc.value.code == "scope_escalation"


def test_an_exchange_must_narrow_the_audience():
    ex, s = exchange_setup()
    with pytest.raises(lab.ExchangeError) as exc:
        ex.exchange(lab.ExchangeRequest(
            subject_token=user_token(s), actor_id="orchestrator",
            audience="agent-platform", scopes=("payments.read",)))
    assert exc.value.code == "no_narrowing"


def test_an_exchange_requires_an_audience():
    ex, s = exchange_setup()
    with pytest.raises(lab.ExchangeError):
        ex.exchange(lab.ExchangeRequest(subject_token=user_token(s),
                                        actor_id="o", scopes=("payments.read",)))


def test_a_non_delegable_token_cannot_be_exchanged():
    ex, s = exchange_setup()
    with pytest.raises(lab.ExchangeError) as exc:
        ex.exchange(lab.ExchangeRequest(
            subject_token=user_token(s, may_delegate=False), actor_id="o",
            audience="investigator", scopes=("payments.read",)))
    assert exc.value.code == "not_delegable"


def test_a_chain_cycle_is_refused():
    ex, s = exchange_setup()
    first = ex.exchange(lab.ExchangeRequest(
        subject_token=user_token(s), actor_id="orchestrator",
        audience="investigator", scopes=("payments.read",)))
    ex2, _ = exchange_setup(audience="investigator")
    with pytest.raises(lab.ExchangeError) as exc:
        ex2.exchange(lab.ExchangeRequest(
            subject_token=first, actor_id="orchestrator",
            audience="core-banking", scopes=("payments.read",)))
    assert exc.value.code == "chain_cycle"


def test_an_agent_cannot_delegate_to_the_user_it_acts_for():
    ex, s = exchange_setup()
    with pytest.raises(lab.ExchangeError) as exc:
        ex.exchange(lab.ExchangeRequest(
            subject_token=user_token(s), actor_id="u-42",
            audience="investigator", scopes=("payments.read",)))
    assert exc.value.code == "chain_cycle"


def test_chain_depth_is_bounded():
    ex, s = exchange_setup(max_depth=1)
    first = ex.exchange(lab.ExchangeRequest(
        subject_token=user_token(s), actor_id="a",
        audience="investigator", scopes=("payments.read",)))
    assert claims_of(first).may_delegate is False       # at the limit, no further hops
    ex2, _ = exchange_setup(audience="investigator", max_depth=1)
    with pytest.raises(lab.ExchangeError) as exc:
        ex2.exchange(lab.ExchangeRequest(
            subject_token=first, actor_id="b",
            audience="core-banking", scopes=("payments.read",)))
    assert exc.value.code == "not_delegable"


def test_impersonation_is_disabled_by_default():
    ex, s = exchange_setup()
    with pytest.raises(lab.ExchangeError) as exc:
        ex.exchange(lab.ExchangeRequest(
            subject_token=user_token(s), actor_id="o", audience="investigator",
            scopes=("payments.read",), delegation=False))
    assert exc.value.code == "impersonation_disabled"


def test_impersonation_erases_the_chain_when_enabled():
    """Which is precisely why a regulated platform disables it."""
    ex, s = exchange_setup(allow_impersonation=True)
    out = ex.exchange(lab.ExchangeRequest(
        subject_token=user_token(s), actor_id="orchestrator", audience="investigator",
        scopes=("payments.read",), delegation=False))
    got = claims_of(out)
    assert got.act == ()
    assert got.sub == "orchestrator"                    # the user is gone


def test_a_derived_token_never_outlives_its_parent():
    ex, s = exchange_setup(now=frozen(1100.0), max_lifetime=1000.0)
    out = ex.exchange(lab.ExchangeRequest(
        subject_token=user_token(s, exp=1150.0), actor_id="o",
        audience="investigator", scopes=("payments.read",), lifetime_seconds=900.0))
    got = claims_of(out)
    assert got.exp <= 1150.0


def test_lifetime_is_capped_by_policy():
    ex, s = exchange_setup(max_lifetime=30.0)
    got = claims_of(ex.exchange(lab.ExchangeRequest(
        subject_token=user_token(s), actor_id="o", audience="investigator",
        scopes=("payments.read",), lifetime_seconds=600.0)))
    assert got.exp - got.iat == 30.0


def test_a_subject_token_alive_only_within_clock_skew_yields_nothing():
    """At now=2020 a token expiring at 1995 still verifies (30s skew) but has NEGATIVE
    remaining lifetime — so a derived credential would be born expired. Refuse."""
    ex, s = exchange_setup(now=frozen(2020.0))
    with pytest.raises(lab.ExchangeError) as exc:
        ex.exchange(lab.ExchangeRequest(
            subject_token=user_token(s, exp=1995.0), actor_id="o",
            audience="investigator", scopes=("payments.read",),
            lifetime_seconds=100.0))
    assert exc.value.code == "expired"


def test_a_nearly_expired_subject_token_yields_a_very_short_credential():
    ex, s = exchange_setup(now=frozen(1990.0))
    got = claims_of(ex.exchange(lab.ExchangeRequest(
        subject_token=user_token(s, exp=1995.0), actor_id="o",
        audience="investigator", scopes=("payments.read",), lifetime_seconds=100.0)))
    assert got.exp - got.iat == 5.0


def test_an_exchange_can_bind_the_result_to_a_key():
    ex, s = exchange_setup()
    got = claims_of(ex.exchange(lab.ExchangeRequest(
        subject_token=user_token(s), actor_id="o", audience="investigator",
        scopes=("payments.read",), bind_to_key="pk")))
    assert got.cnf == lab.key_thumbprint("pk")


def test_the_tenant_survives_every_hop():
    ex, s = exchange_setup()
    got = claims_of(ex.exchange(lab.ExchangeRequest(
        subject_token=user_token(s), actor_id="o", audience="investigator",
        scopes=("payments.read",))))
    assert got.tenant == "wholesale"


# ======================================================================================
# 7. SPIFFE / SPIRE
# ======================================================================================


def test_spiffe_id_parsing_and_rendering():
    parsed = lab.SpiffeID.parse("spiffe://bank.ae/ns/agents/sa/investigator")
    assert parsed.trust_domain == "bank.ae"
    assert parsed.path == "/ns/agents/sa/investigator"
    assert str(parsed) == "spiffe://bank.ae/ns/agents/sa/investigator"


@pytest.mark.parametrize("bad", ["http://bank.ae/x", "spiffe://bank.ae", "spiffe://", "x"])
def test_invalid_spiffe_ids_are_refused(bad):
    with pytest.raises(ValueError):
        lab.SpiffeID.parse(bad)


def spire():
    server = lab.SpireServer(trust_domain="bank.ae", now=frozen(1000.0), svid_ttl=60.0)
    server.register(lab.RegistrationEntry(
        "spiffe://bank.ae/ns/agents/sa/investigator",
        (lab.Selector("k8s:ns", "agents"), lab.Selector("k8s:sa", "investigator"))))
    return server


def test_attestation_issues_an_svid_from_verified_properties():
    svid = spire().attest([lab.Selector("k8s:ns", "agents"),
                           lab.Selector("k8s:sa", "investigator"),
                           lab.Selector("k8s:pod", "abc")], public_key="pk")
    assert str(svid.spiffe_id) == "spiffe://bank.ae/ns/agents/sa/investigator"
    assert svid.is_valid(1000.0)
    assert not svid.is_valid(1100.0)


def test_all_registered_selectors_must_match():
    """A subset match would let a workload with one matching property claim an identity
    meant for a narrower set."""
    with pytest.raises(lab.TokenError) as exc:
        spire().attest([lab.Selector("k8s:ns", "agents")], public_key="pk")
    assert exc.value.code == "attestation_failed"


def test_ambiguous_registration_is_refused_rather_than_guessed():
    server = spire()
    server.register(lab.RegistrationEntry(
        "spiffe://bank.ae/ns/agents/sa/other",
        (lab.Selector("k8s:ns", "agents"),)))
    with pytest.raises(lab.TokenError) as exc:
        server.attest([lab.Selector("k8s:ns", "agents"),
                       lab.Selector("k8s:sa", "investigator")], public_key="pk")
    assert exc.value.code == "ambiguous_attestation"


def test_a_registration_entry_needs_selectors():
    with pytest.raises(ValueError):
        spire().register(lab.RegistrationEntry("spiffe://bank.ae/ns/x/sa/y", ()))


def test_an_entry_outside_the_trust_domain_is_refused():
    with pytest.raises(ValueError):
        spire().register(lab.RegistrationEntry(
            "spiffe://other.example/ns/x/sa/y", (lab.Selector("a", "b"),)))


def test_svid_serials_increase():
    server = spire()
    first = server.attest([lab.Selector("k8s:ns", "agents"),
                           lab.Selector("k8s:sa", "investigator")], public_key="pk")
    second = server.attest([lab.Selector("k8s:ns", "agents"),
                            lab.Selector("k8s:sa", "investigator")], public_key="pk")
    assert second.serial > first.serial


def test_mtls_allows_a_registered_caller():
    svid = spire().attest([lab.Selector("k8s:ns", "agents"),
                           lab.Selector("k8s:sa", "investigator")], public_key="pk")
    lab.mtls_authorize(svid, lab.MtlsPolicy(
        allowed_callers=("spiffe://bank.ae/ns/agents/sa/investigator",)), now=1000.0)


def test_mtls_refuses_an_expired_svid():
    svid = spire().attest([lab.Selector("k8s:ns", "agents"),
                           lab.Selector("k8s:sa", "investigator")], public_key="pk")
    with pytest.raises(lab.TokenError) as exc:
        lab.mtls_authorize(svid, lab.MtlsPolicy(
            allowed_callers=(str(svid.spiffe_id),)), now=99999.0)
    assert exc.value.code == "svid_expired"


def test_mtls_refuses_an_unlisted_caller():
    svid = spire().attest([lab.Selector("k8s:ns", "agents"),
                           lab.Selector("k8s:sa", "investigator")], public_key="pk")
    with pytest.raises(lab.TokenError):
        lab.mtls_authorize(svid, lab.MtlsPolicy(allowed_callers=()), now=1000.0)


def test_cross_domain_requires_explicit_federation():
    foreign = lab.SVID(lab.SpiffeID.parse("spiffe://partner.example/ns/x/sa/y"),
                       expires_at=1e12, serial=1, public_key="pk")
    with pytest.raises(lab.TokenError):
        lab.mtls_authorize(foreign, lab.MtlsPolicy(allowed_callers=()), now=1000.0)
    lab.mtls_authorize(foreign, lab.MtlsPolicy(
        allowed_callers=("spiffe://partner.example/ns/x/sa/y",),
        federated_domains=("partner.example",)), now=1000.0)


# ======================================================================================
# 8. NHI lifecycle and JIT credentials
# ======================================================================================


def identity(**kwargs):
    kwargs.setdefault("identity_id", "investigator")
    kwargs.setdefault("kind", "agent")
    kwargs.setdefault("owner", "layla")
    kwargs.setdefault("tenant", "wholesale")
    kwargs.setdefault("allowed_scopes", ("payments.read", "crm.read"))
    return lab.NonHumanIdentity(**kwargs)


def registry(active=True):
    reg = lab.IdentityRegistry(now=frozen(1000.0))
    reg.register(identity())
    if active:
        reg.transition("investigator", lab.NHIState.APPROVED)
        reg.transition("investigator", lab.NHIState.ACTIVE)
    return reg


def test_every_nhi_needs_a_human_owner():
    with pytest.raises(ValueError):
        lab.IdentityRegistry(now=frozen()).register(identity(owner=""))


def test_duplicate_registration_is_refused():
    with pytest.raises(ValueError):
        registry().register(identity())


def test_the_lifecycle_is_a_declared_state_machine():
    reg = lab.IdentityRegistry(now=frozen())
    reg.register(identity())
    with pytest.raises(ValueError):
        reg.transition("investigator", lab.NHIState.ACTIVE)      # must be approved first
    reg.transition("investigator", lab.NHIState.APPROVED)
    reg.transition("investigator", lab.NHIState.ACTIVE)
    reg.transition("investigator", lab.NHIState.SUSPENDED)
    reg.transition("investigator", lab.NHIState.ACTIVE)
    reg.transition("investigator", lab.NHIState.RETIRED)
    with pytest.raises(ValueError):
        reg.transition("investigator", lab.NHIState.ACTIVE)      # retired is terminal


def test_only_an_active_identity_may_receive_credentials():
    reg = registry()
    assert reg.may_receive_credentials("investigator")
    reg.transition("investigator", lab.NHIState.SUSPENDED)
    assert not reg.may_receive_credentials("investigator")


def test_an_unknown_identity_raises():
    with pytest.raises(KeyError):
        registry().get("ghost")


def broker(reg=None, now=None):
    now = now or frozen(1100.0)
    s = signer()
    v = lab.Verifier(signers={"k1": s},
                     policy=lab.VerificationPolicy(
                         audience="agent-platform",
                         trusted_issuers=("https://login.bank.ae", "https://sts.bank.ae")),
                     now=now)
    ex = lab.TokenExchange(issuer="https://sts.bank.ae", signer=s, verifier=v, now=now)
    return lab.JitCredentialBroker(registry=reg or registry(), exchange=ex,
                                   issuer="https://sts.bank.ae", signer=s, now=now), s


def test_a_workload_credential_is_short_lived_and_narrowly_audienced():
    b, _ = broker()
    got = claims_of(b.issue(lab.CredentialRequest(
        identity_id="investigator", audience="crm", scopes=("crm.read",),
        lifetime_seconds=60.0)))
    assert got.aud == "crm"
    assert got.scope == ("crm.read",)
    assert got.exp - got.iat == 60.0
    assert got.may_delegate is False


def test_a_workload_credential_records_the_agent_as_the_actor():
    b, _ = broker()
    got = claims_of(b.issue(lab.CredentialRequest("investigator", "crm", ("crm.read",))))
    assert [a.sub for a in got.act] == ["investigator"]


def test_a_credential_cannot_exceed_the_agents_registered_scopes():
    b, _ = broker()
    with pytest.raises(lab.ScopeEscalation):
        b.issue(lab.CredentialRequest("investigator", "crm", ("payments.release",)))


def test_lifetime_is_capped_by_the_identity_and_the_broker():
    reg = lab.IdentityRegistry(now=frozen(1000.0))
    reg.register(identity(max_credential_lifetime=15.0))
    reg.transition("investigator", lab.NHIState.APPROVED)
    reg.transition("investigator", lab.NHIState.ACTIVE)
    b, _ = broker(reg=reg)
    got = claims_of(b.issue(lab.CredentialRequest(
        "investigator", "crm", ("crm.read",), lifetime_seconds=600.0)))
    assert got.exp - got.iat == 15.0


def test_a_suspended_identity_cannot_receive_a_credential():
    reg = registry()
    reg.transition("investigator", lab.NHIState.SUSPENDED)
    b, _ = broker(reg=reg)
    with pytest.raises(lab.TokenError) as exc:
        b.issue(lab.CredentialRequest("investigator", "crm", ("crm.read",)))
    assert exc.value.code == "identity_not_active"


def test_acting_for_a_user_exchanges_and_keeps_the_user_in_the_chain():
    b, s = broker()
    subject = user_token(s, scope=("crm.read",))
    got = claims_of(b.issue(lab.CredentialRequest(
        identity_id="investigator", audience="crm", scopes=("crm.read",),
        subject_token=subject)))
    assert got.sub == "u-42"
    assert [a.sub for a in got.act] == ["investigator"]
    assert lab.describe_chain(got) == "u-42 -> investigator"


def test_a_bound_jit_credential_carries_a_confirmation_claim():
    b, _ = broker()
    got = claims_of(b.issue(lab.CredentialRequest(
        "investigator", "crm", ("crm.read",), bind_to_key="pk")))
    assert got.cnf == lab.key_thumbprint("pk")


# ======================================================================================
# 9. Determinism
# ======================================================================================


def test_two_identical_flows_produce_identical_tokens():
    a = auth_server(now=clock())
    b = auth_server(now=clock())
    args = dict(client_id="web", redirect_uri="https://bank.ae/cb",
                scopes=["payments.read"], challenge=lab.code_challenge(VERIFIER_SECRET),
                user_id="u-42", tenant="wholesale")
    ca, cb = a.authorize(**args), b.authorize(**args)
    ta = a.exchange_code(code=ca, client_id="web", verifier=VERIFIER_SECRET,
                         redirect_uri="https://bank.ae/cb", audience="agent-platform")
    tb = b.exchange_code(code=cb, client_id="web", verifier=VERIFIER_SECRET,
                         redirect_uri="https://bank.ae/cb", audience="agent-platform")
    assert ta == tb