import secrets

import pytest
from cryptography.hazmat.primitives.asymmetric.mlkem import MLKEM768PrivateKey
from cryptography.hazmat.primitives.asymmetric.x25519 import X25519PublicKey

from quantumpipe.chab3 import encrypt
from quantumpipe.exceptions import ProtocolError
from quantumpipe.handshake import (
    _MSG1_REGISTER,
    _fresh_ephemeral,
    _public_from_idroot,
    _x25519_private_key,
)
from quantumpipe.kdf import kdf
from quantumpipe.sansio import (
    _CREDENTIAL,
    Established,
    Message,
    Session,
    credential_id,
    generate_registration_secret,
)


def make_sessions(cred_a, cred_b=None, ad=b"") -> tuple[Session, Session]:
    """A Session pair over fresh key stores; cred_b defaults to cred_a."""
    a = Session({"x": cred_a}, "x", ad=ad)
    b = Session({"x": cred_a if cred_b is None else cred_b}, ad=ad)
    return a, b


def run_handshake(
    a: Session, b: Session, answer=True
) -> tuple[Established, Established]:
    (m1,) = a()
    (m2,) = b(m1)
    out = a(m2)
    if isinstance(out[-1], bytes):  # registration: two gated tail messages
        (m3,) = out
        (vb,) = b(m3)  # the responder gates first
        vb.answer = answer
        (m4,) = b(vb)  # his answer releases msg4
        (va,) = a(m4)
        assert va.sas == vb.sas
        va.answer = answer
        m5, est_a = a(va)  # the initiator's answer finishes her side
        (est_b,) = b(m5)  # and msg5 finishes the responder's
        assert est_a.sas == est_b.sas == va.sas
        assert est_a.verified == est_b.verified == (answer is True)
    else:  # verified credential: no comparison, established right away
        m3, est_a = out
        (est_b,) = b(m3)
        assert est_a.sas == est_b.sas  # the stored digest
        assert est_a.verified and est_b.verified
    assert isinstance(est_a, Established)
    assert a.state == b.state == "established"
    # both ends derived matching channels: traffic roundtrips both ways
    (w,) = a(Message(b"ping", 1))
    assert b(w) == [Message(b"ping", 1)]
    (w,) = b(Message(b"pong", 1))
    assert a(w) == [Message(b"pong", 1)]
    return est_a, est_b


@pytest.fixture
def psk():
    # both parties simply store the registration secret itself
    return generate_registration_secret()


@pytest.fixture
def registered(psk):
    a, b = make_sessions(psk)
    run_handshake(a, b)
    return a.credential, b.credential


def test_registration_credential(psk):
    a, b = make_sessions(psk)
    run_handshake(a, b)
    blob_a, blob_b = a.credential, b.credential
    assert len(blob_a) == len(blob_b) == _CREDENTIAL.size
    id_a, shared_a, pub_b_m, pub_b_x, sas_a, vf_a = _CREDENTIAL.unpack(blob_a)
    id_b, shared_b, pub_a_m, pub_a_x, sas_b, vf_b = _CREDENTIAL.unpack(blob_b)
    assert id_a != id_b  # fresh random identity roots per party
    assert shared_a == shared_b  # same rotated shared key
    assert (pub_a_m, pub_a_x) == _public_from_idroot(id_a)  # pubs match idroots
    assert (pub_b_m, pub_b_x) == _public_from_idroot(id_b)
    assert sas_a == sas_b  # the stored comparison digest
    assert vf_a and vf_b  # confirmed by run_handshake's answers


def test_key_change_is_atomic(psk):
    # the composite replaces the registration secret only at the
    # key-change points -- the initiator's swap at her answer, the
    # responder's when msg5 completes; until then the registration
    # secret keeps serving new connections
    a, b = make_sessions(psk)
    (m1,) = a()
    (m2,) = b(m1)
    (m3,) = a(m2)
    assert a._store["x"] == psk
    (vb,) = b(m3)
    assert b._store["x"] == psk
    (m4,) = b(vb)  # answered (default None), but bob's swap waits for msg5
    assert b._store["x"] == psk
    (va,) = a(m4)
    m5, _est_a = a(va)  # alice's swap happens at her answer
    assert len(a._store["x"]) == _CREDENTIAL.size
    assert b._store["x"] == psk  # bob's still pending
    b(m5)  # msg5 completes bob's side: swap
    assert len(b._store["x"]) == _CREDENTIAL.size


def test_unanswered_registration_retries(psk):
    # a connection dropped mid-registration leaves the registration
    # secret live; the next connection registers afresh
    store_a, store_b = {"x": psk}, {"x": psk}
    a, b = Session(store_a, "x"), Session(store_b)
    (m1,) = a()
    (m2,) = b(m1)
    a(m2)  # parked awaiting msg4, never answered
    assert store_a["x"] == psk and store_b["x"] == psk
    run_handshake(Session(store_a, "x"), Session(store_b))  # simply retried


def test_traffic_roundtrip(psk):
    a, b = make_sessions(psk)
    run_handshake(a, b)
    (w,) = a(Message(b"one", 1))
    assert b(w) == [Message(b"one", 1)]
    (w,) = a(Message(b"", 1))
    assert b(w) == [Message(b"", 1)]  # empty payloads work


def test_user_kinds_pass_through(psk):
    # the kind is opaque to the state machine: any nonzero value roundtrips
    a, b = make_sessions(psk)
    run_handshake(a, b)
    for kind in (1, 2, 7, 255):
        (w,) = a(Message(b"payload", kind))
        assert b(w) == [Message(b"payload", kind)]


def test_protocol_kind_reserved(psk):
    # kind 0 is reserved for in-band protocol messages: rejected both ways
    a, b = make_sessions(psk)
    run_handshake(a, b)
    with pytest.raises(ProtocolError, match="reserved"):
        a(Message(b"x", 0))
    w = bytes(a._tx.encrypt(bytes([0]) + b"x"))  # inbound kind 0
    with pytest.raises(ProtocolError, match="reserved"):
        b(w)
    with pytest.raises(ProtocolError):
        a(Message(b"x", 300))  # does not fit the kind byte


def test_traffic_tamper_rejected(psk):
    # any tampering with a traffic frame fails the open; failed opens
    # consume no nonce, so the untampered frame still opens afterwards
    a, b = make_sessions(psk)
    run_handshake(a, b)
    cred = b.credential
    (w,) = a(Message(b"ping", 1))
    w = bytearray(w)
    orig = bytes(w)
    w[-1] ^= 0x01
    with pytest.raises(ProtocolError):
        b(bytes(w))
    w[0] ^= 0x01  # a different spot: still a tag failure
    with pytest.raises(ProtocolError):
        b(bytes(w))
    assert b.credential == cred  # an AEAD failure burns nothing
    assert b(orig) == [Message(b"ping", 1)]  # no nonce consumed, still in sync


def test_traffic_in_order(psk):
    # the channel is sequential: out-of-order frames fail without
    # consuming a nonce
    a, b = make_sessions(psk)
    run_handshake(a, b)
    (w1,) = a(Message(b"one", 1))
    (w2,) = a(Message(b"two", 1))
    with pytest.raises(ProtocolError):
        b(w2)  # out of order
    assert b(w1) == [Message(b"one", 1)]
    assert b(w2) == [Message(b"two", 1)]


def test_traffic_before_establishment_rejected(psk):
    a, b = make_sessions(psk)
    (m1,) = a()
    (m2,) = b(m1)
    (_m3,) = a(m2)  # initiator parked awaiting msg4
    with pytest.raises(ProtocolError, match="no traffic"):
        a(Message(b"early", 1))
    with pytest.raises(ProtocolError, match="no traffic"):
        b(Message(b"early", 1))  # responder still awaiting msg3


def test_resume_does_not_rotate(registered):
    blob_a, blob_b = registered
    a, b = make_sessions(blob_a, blob_b)
    run_handshake(a, b)
    assert a.credential == blob_a  # no rotation in ordinary operation
    assert b.credential == blob_b


def test_resume_fresh_session_keys(registered):
    blob_a, blob_b = registered
    a1, b1 = make_sessions(blob_a, blob_b)
    run_handshake(a1, b1)
    (ct1,) = a1(Message(b"same plaintext", 1))
    a2, b2 = make_sessions(blob_a, blob_b)
    run_handshake(a2, b2)
    (ct2,) = a2(Message(b"same plaintext", 1))
    assert ct2 != ct1  # fresh ephemeral keys


def test_role_reversal(registered):
    a, b = make_sessions(registered[1], registered[0])  # either may initiate
    run_handshake(a, b)


def test_state_strings(psk):
    a, b = make_sessions(psk)
    assert a.state == "reg:hello"
    assert b.state == "hello"  # mode unknown until msg1 selects a key
    (m1,) = a()
    assert a.state == "reg:msg2"
    (m2,) = b(m1)
    assert b.state == "reg:msg3"
    (m3,) = a(m2)
    assert a.state == "reg:msg4"  # msg3 out, awaiting bob's challenge
    (vb,) = b(m3)
    assert b.state == "reg:verify"  # the responder gates first
    (m4,) = b(vb)  # default answer: proceed unverified
    assert b.state == "reg:msg5"
    (va,) = a(m4)
    assert a.state == "reg:verify"
    (m5, _est_a) = a(va)
    assert a.state == "established"
    (_est_b,) = b(m5)
    assert b.state == "established"


def test_established_gated_on_answer(psk):
    # the comparison answers gate the tail of the exchange: the
    # responder's answer releases msg4, the initiator's releases msg5
    # and finishes her side; msg5 finishes the responder's
    a, b = make_sessions(psk)
    (m1,) = a()
    (m2,) = b(m1)
    (m3,) = a(m2)
    (vb,) = b(m3)
    assert a.state == "reg:msg4" and b.state == "reg:verify"
    with pytest.raises(ProtocolError):
        _ = b.verified  # handshake not complete
    vb.answer = True
    (m4,) = b(vb)
    assert b.state == "reg:msg5"  # answered, but not established
    with pytest.raises(ProtocolError):
        _ = b.verified  # still not complete
    (va,) = a(m4)
    va.answer = True
    m5, est_a = a(va)
    assert est_a.verified and est_a.sas == va.sas
    (est_b,) = b(m5)
    assert est_b.verified and est_b.sas == va.sas


def test_verification_reject(psk):
    a, b = make_sessions(psk)
    (m1,) = a()
    (m2,) = b(m1)
    (m3,) = a(m2)
    (vb,) = b(m3)
    (m4,) = b(vb)  # responder proceeds (default answer)
    (va,) = a(m4)
    va.answer = False
    with pytest.raises(ProtocolError, match="rejected"):
        a(va)
    assert a.state == "failed"
    assert a._store == {}  # rejected: the credential is deleted; no msg5
    with pytest.raises(ProtocolError):
        a()


def test_verification_reject_responder(psk):
    a, b = make_sessions(psk)
    (m1,) = a()
    (m2,) = b(m1)
    (m3,) = a(m2)
    (vb,) = b(m3)
    vb.answer = False
    with pytest.raises(ProtocolError, match="rejected"):
        b(vb)
    assert b.state == "failed"
    assert b._store == {}  # rejected: the credential is deleted; no msg4


def test_verification_not_required(psk):
    # answer=None (the default) skips the comparison and continues,
    # flagged unverified
    a, b = make_sessions(psk)
    (m1,) = a()
    (m2,) = b(m1)
    (m3,) = a(m2)
    (vb,) = b(m3)
    (m4,) = b(vb)
    (va,) = a(m4)
    m5, est_a = a(va)
    (est_b,) = b(m5)
    assert isinstance(est_a, Established) and isinstance(est_b, Established)
    assert not est_a.verified and not est_b.verified
    assert not a.verified and not b.verified
    assert est_a.sas == est_b.sas  # the SAS is still exposed
    assert a.state == b.state == "established"
    assert _CREDENTIAL.unpack(a._store["x"])[5] is False  # stored unverified


def test_wrong_calls_rejected(psk):
    a, b = make_sessions(psk)
    with pytest.raises(ProtocolError):
        _ = a.verification  # too early
    with pytest.raises(ProtocolError):
        _ = a.verified  # not complete
    (m1,) = a()
    with pytest.raises(ProtocolError):
        a()  # local advance not possible while expecting msg2
    (m2,) = b(m1)
    with pytest.raises(ProtocolError):
        b()  # msg3 expected, local advance not possible
    (m3,) = a(m2)
    with pytest.raises(ProtocolError):
        a()  # msg4 expected, local advance not possible
    with pytest.raises(ProtocolError):
        a(b"junk")  # msg4 expected, junk does not open
    (vb,) = b(m3)
    with pytest.raises(ProtocolError):
        b(b"junk")  # verify takes only the comparison answer
    (m4,) = b(vb)
    (_v,) = a(m4)
    with pytest.raises(ProtocolError):
        a(b"junk")  # verify takes only the comparison answer
    with pytest.raises(ProtocolError, match="key_id"):
        Session({})()  # initiation requires a key_id


def test_session_ad(psk, registered):
    ad = b"myapp v2"
    for cred_a, cred_b in ((psk, psk), registered):
        a, b = make_sessions(cred_a, cred_b, ad=ad)
        run_handshake(a, b)


def test_mismatched_ad_rejected(psk, registered):
    for cred_a, cred_b in ((psk, psk), registered):
        a = Session({"x": cred_a}, "x", ad=b"v1")
        b = Session({"x": cred_b}, ad=b"v2")
        (m1,) = a()
        with pytest.raises(ProtocolError):
            b(m1)
        assert b._store["x"] == cred_b  # an AEAD failure burns nothing


def test_tampered_msg1_rejected(psk, registered):
    for cred_a, cred_b in ((psk, psk), registered):
        a, b = make_sessions(cred_a, cred_b)
        (m1,) = a()
        m1 = bytearray(m1)
        m1[-1] ^= 0x01
        with pytest.raises(ProtocolError):
            b(bytes(m1))
        assert b._store["x"] == cred_b  # an AEAD failure burns nothing


def test_tampered_msg2_rejected(psk, registered):
    for cred_a, cred_b in ((psk, psk), registered):
        a, b = make_sessions(cred_a, cred_b)
        (m1,) = a()
        m2 = bytearray(b(m1)[0])
        m2[-1] ^= 0x01
        with pytest.raises(ProtocolError):
            a(bytes(m2))
        assert a._store["x"] == cred_a  # an AEAD failure burns nothing


def test_tampered_msg3_rejected(psk, registered):
    for cred_a, cred_b in ((psk, psk), registered):
        a, b = make_sessions(cred_a, cred_b)
        (m1,) = a()
        m2 = b(m1)[0]
        m3 = bytearray(a(m2)[0])  # msg3 comes first in both flows
        m3[-1] ^= 0x01
        with pytest.raises(ProtocolError):
            b(bytes(m3))
        assert b._store["x"] == cred_b  # an AEAD failure burns nothing


def test_tampered_msg4_rejected(psk):
    a, b = make_sessions(psk)
    (m1,) = a()
    (m2,) = b(m1)
    (m3,) = a(m2)
    (vb,) = b(m3)
    m4 = bytearray(b(vb)[0])
    m4[-1] ^= 0x01
    with pytest.raises(ProtocolError):
        a(bytes(m4))
    assert a._store["x"] == psk  # an AEAD failure burns nothing


def test_tampered_msg5_rejected(psk):
    a, b = make_sessions(psk)
    (m1,) = a()
    (m2,) = b(m1)
    (m3,) = a(m2)
    (vb,) = b(m3)
    (m4,) = b(vb)
    (va,) = a(m4)
    m5 = bytearray(a(va)[0])  # registration's last message
    m5[-1] ^= 0x01
    with pytest.raises(ProtocolError):
        b(bytes(m5))
    assert b._store["x"] == psk  # an AEAD failure burns nothing


def test_mismatched_shared_rejected(psk, registered):
    # bob still holds the bare registration secret, alice the rotated
    # composite: the shared keys differ, so her msg1 box does not open
    a = Session({"x": registered[0]}, "x")
    b = Session({"x": psk})
    (m1,) = a()
    with pytest.raises(ProtocolError):
        b(m1)
    assert b._store["x"] == psk  # no burn


def test_forged_idroot_rejected(registered):
    # resume composite with a wrong idroot: decapsulation yields garbage,
    # so bob's tag fails to verify -- invalid data inside a valid box,
    # and the forged credential burns
    blob = registered[0]
    bad = _CREDENTIAL.pack(secrets.token_bytes(32), *_CREDENTIAL.unpack(blob)[1:])
    a, b = make_sessions(bad, registered[1])
    (m1,) = a()
    (m2,) = b(m1)
    with pytest.raises(ProtocolError, match="burned"):
        a(m2)
    assert a._store == {}  # burned
    assert b._store["x"] == registered[1]  # bob's side untouched


def test_verified_resumption_is_silent(registered):
    # a verified credential authenticates via the pinned statics: no
    # comparison is asked
    a, b = make_sessions(registered[0], registered[1])
    (m1,) = a()
    (m2,) = b(m1)  # no Verification emitted
    assert b.state == "auth:msg3"
    with pytest.raises(ProtocolError):
        _ = b.verification  # nothing to verify
    m3, est_a = a(m2)
    assert isinstance(est_a, Established)
    assert est_a.verified and est_a.sas  # pinned statics, stored digest
    assert a.state == "established"
    (est_b,) = b(m3)
    assert isinstance(est_b, Established)
    assert est_b.verified and est_b.sas == est_a.sas


def test_unverified_resumption_is_silent(psk):
    # registration skipped: the stored composite is unverified, but the
    # resumption flow never prompts -- the stored digest rides on
    # Established and the parties confirm or reject whenever they wish
    a, b = make_sessions(psk)
    run_handshake(a, b, answer=None)
    cred_a, cred_b = a.credential, b.credential
    assert _CREDENTIAL.unpack(cred_a)[5] is False

    a2, b2 = make_sessions(cred_a, cred_b)
    (m1,) = a2()
    (m2,) = b2(m1)  # no Verification emitted on the resumption flow
    m3, est_a = a2(m2)
    (est_b,) = b2(m3)
    assert a2.state == b2.state == "established"
    assert not est_a.verified and not est_b.verified
    # the stored digest is still exposed for an optional comparison
    assert est_a.sas == est_b.sas == a2.verification.sas
    a2.confirm(True)  # compared out of band: flag it verified
    b2.confirm(True)
    assert a2.verified and b2.verified
    assert _CREDENTIAL.unpack(a2._store["x"])[5] is True  # flag flipped

    a3, b3 = make_sessions(a2._store["x"], b2._store["x"])
    run_handshake(a3, b3)  # verified: silent as before


def test_resumption_reject_deletes(psk):
    a, b = make_sessions(psk)
    run_handshake(a, b, answer=None)
    a2, b2 = make_sessions(a.credential, b.credential)
    (m1,) = a2()
    (m2,) = b2(m1)
    m3, _est_a = a2(m2)
    b2(m3)
    a2.confirm(False)  # rejected even on the resumption flow: gone
    assert a2._store == {}


def test_responder_scans_key_store(psk):
    # a server holds credentials for many callers; msg1 picks the one
    # that opens it, and only that entry is touched
    other1, other2 = generate_registration_secret(), generate_registration_secret()
    a = Session({"alice": psk}, "alice")
    b = Session({"alice": psk, "bob": other1, "carol": other2})
    run_handshake(a, b)
    assert len(a._store["alice"]) == _CREDENTIAL.size
    assert len(b._store["alice"]) == _CREDENTIAL.size  # registered
    assert b._store["bob"] == other1  # untouched
    assert b._store["carol"] == other2
    # same rotated key on both sides
    assert (
        _CREDENTIAL.unpack(a._store["alice"])[1]
        == _CREDENTIAL.unpack(b._store["alice"])[1]
    )


def test_unknown_caller_no_burn(psk):
    stranger = Session({"x": generate_registration_secret()}, "x")
    b = Session({"alice": psk})
    (m1,) = stranger()
    with pytest.raises(ProtocolError):
        b(m1)  # no credential opens the box
    assert b._store == {"alice": psk}  # nothing burned


def test_parallel_registration_loser_fails(psk):
    # two registrations race on one registration key: the first exchange
    # to complete msg5 swaps the slot -- the loser's finish finds the
    # registration key gone and fails instead of overwriting
    srv = {"alice": psk}
    x, y = Session({"s": psk}, "s"), Session({"s": psk}, "s")
    sx, sy = Session(srv), Session(srv)  # two responder sessions, one store
    (m1x,) = x()
    (m1y,) = y()
    (m2x,) = sx(m1x)
    (m2y,) = sy(m1y)  # both in flight on the same key
    (m3x,) = x(m2x)
    (m3y,) = y(m2y)
    (vx,) = sx(m3x)
    (vy,) = sy(m3y)  # both responders parked; the key is still live
    assert srv["alice"] == psk
    (m4x,) = sx(vx)  # default answers release the msg4s
    (m4y,) = sy(vy)
    (vx2,) = x(m4x)
    (vy2,) = y(m4y)
    m5x, _ex = x(vx2)
    (est_x,) = sx(m5x)  # X completes first: wins the swap
    assert isinstance(est_x, Established)
    assert len(srv["alice"]) == _CREDENTIAL.size
    m5y, _ey = y(vy2)
    with pytest.raises(ProtocolError, match="spent"):
        sy(m5y)  # Y finds the registration key gone
    assert sy.state == "failed"
    assert len(srv["alice"]) == _CREDENTIAL.size  # the winner is intact


def test_parallel_registration_same_client_store(psk):
    # the same arbitration on the initiator's store: the second session
    # to answer loses
    cli = {"s": psk}
    x, y = Session(cli, "s"), Session(cli, "s")
    b1, b2 = Session({"s": psk}), Session({"s": psk})
    (m1x,) = x()
    (m1y,) = y()
    (m2x,) = b1(m1x)
    (m2y,) = b2(m1y)
    (m3x,) = x(m2x)
    (m3y,) = y(m2y)
    (v1,) = b1(m3x)
    (v2,) = b2(m3y)
    (m4x,) = b1(v1)  # default answers release the msg4s
    (m4y,) = b2(v2)
    (vx,) = x(m4x)
    (vy,) = y(m4y)  # both parked; the registration key is still live
    assert cli["s"] == psk
    vx.answer = None  # even a skip swaps: X wins
    _m5x, est_x = x(vx)
    assert isinstance(est_x, Established)
    assert len(cli["s"]) == _CREDENTIAL.size
    vy.answer = None
    with pytest.raises(ProtocolError, match="spent"):
        y(vy)
    assert len(cli["s"]) == _CREDENTIAL.size  # not overwritten


def test_loser_burn_spares_winner(psk):
    # the loser's content failure must not burn the winner's freshly
    # swapped credential
    srv = {"alice": psk}
    x, y = Session({"s": psk}, "s"), Session({"s": psk}, "s")
    sx, sy = Session(srv), Session(srv)
    (m1x,) = x()
    (m1y,) = y()
    (m2x,) = sx(m1x)
    (m2y,) = sy(m1y)
    (m3x,) = x(m2x)
    (vx,) = sx(m3x)
    (m4x,) = sx(vx)
    (vx2,) = x(m4x)
    m5x, _ex = x(vx2)
    sx(m5x)  # X completes first: wins the swap
    y._exchange._reveal = secrets.token_bytes(32)  # Y lies about its reveal
    (m3y,) = y(m2y)
    with pytest.raises(ProtocolError, match="commitment"):
        sy(m3y)
    assert len(srv["alice"]) == _CREDENTIAL.size  # winner intact, no burn


def test_flag_flip_tolerates_parallel_confirm(psk):
    # two resumptions of the same credential confirm it concurrently;
    # the second flag flip must recognize it as the same credential
    # (only the flag byte changed underneath it)
    cli, srv = {"s": psk}, {"s": psk}
    run_handshake(Session(cli, "s"), Session(srv), answer=None)
    r1, r2 = Session(cli, "s"), Session(cli, "s")  # same store, two connections
    s1, s2 = Session(srv), Session(srv)
    (m1_1,) = r1()
    (m1_2,) = r2()
    (m2_1,) = s1(m1_1)
    (m2_2,) = s2(m1_2)
    m3_1, _e1 = r1(m2_1)
    m3_2, _e2 = r2(m2_2)
    s1(m3_1)
    s2(m3_2)
    r1.confirm(True)  # flips the slot to verified
    r2.confirm(True)  # slot changed under it, but same credential
    assert r1.verified and r2.verified
    assert _CREDENTIAL.unpack(cli["s"])[5] is True


def test_commitment_reveal_checked(psk):
    # a reveal that does not open the msg1 commitment is invalid data
    # inside a valid box: the credential burns
    a, b = make_sessions(psk)
    (m1,) = a()
    (m2,) = b(m1)
    a._exchange._reveal = secrets.token_bytes(32)  # alice lies about her reveal
    (m3,) = a(m2)
    with pytest.raises(ProtocolError, match="commitment"):
        b(m3)
    assert b._store == {}  # burned


def test_mitm_digest_not_steerable(psk):
    # A MITM who knows the registration secret completes independent legs
    # with both parties. Bob's digest derives from the commitment the MITM
    # made in msg1 -- before alice's reveal could be known -- so it cannot
    # be steered toward alice's digest: the two values simply differ.
    a = Session({"x": psk}, "x")
    m_in = Session({"x": psk})
    m_out = Session({"x": psk}, "x")
    b = Session({"x": psk})
    (m1,) = a()
    (m2m,) = m_in(m1)  # the MITM answers alice
    (m3a,) = a(m2m)
    (vm,) = m_in(m3a)
    (m4m,) = m_in(vm)
    (va,) = a(m4m)  # alice's digest
    (m1m,) = m_out()  # and opens his own leg toward bob
    (m2b,) = b(m1m)
    (m3m,) = m_out(m2b)
    (vb,) = b(m3m)  # bob's digest
    assert va.sas != vb.sas


def test_mitm_reused_reveal_detected(psk):
    # A MITM who finishes a leg with alice and steals her reveal cannot
    # reuse it toward bob: the sas covers the per-connection ee, which
    # the responder contributes only after the commitment -- so a
    # recommitted reveal still produces a different digest.
    a = Session({"x": psk}, "x")
    m_in = Session({"x": psk})
    m_out = Session({"x": psk}, "x")
    b = Session({"x": psk})
    # leg alpha: the MITM completes with alice and steals her reveal
    (m1,) = a()
    (m2m,) = m_in(m1)
    (m3a,) = a(m2m)
    (vm,) = m_in(m3a)
    r_a = a._exchange._reveal
    # leg beta: the MITM recommits the stolen reveal toward bob
    hs = m_out._exchange
    hs._reveal = r_a
    hs._eph_m, hs._ephemeral, e_m, e_x = _fresh_ephemeral()
    body = _MSG1_REGISTER.pack(e_m, e_x, kdf("commit", e_m, e_x, r_a))
    msg1m = bytes(encrypt(body, key=kdf("reg/msg1", psk), ad=b""))
    hs._chain.advance(msg1m)  # startA was bypassed: advance the chain
    (m2b,) = b(msg1m)
    m3m = hs.respondA(m2b)
    (vb,) = b(m3m)  # bob's checks pass: the commitment is valid...
    (m4m,) = m_in(vm)
    (va,) = a(m4m)
    assert va.sas != vb.sas  # ...but the digests differ


def test_relayed_reveal_rejected(psk):
    # A MITM relaying alice's msg1 verbatim cannot reuse her msg3 either:
    # the reveal box is keyed and ad-chained to her own leg only.
    a = Session({"x": psk}, "x")
    m = Session({"x": psk})
    b = Session({"x": psk})
    (m1,) = a()
    b(m1)  # bob answers the relayed hello
    (m2m,) = m(m1)  # the MITM answers alice too
    (m3,) = a(m2m)  # alice reveals r -- on the MITM's leg only
    with pytest.raises(ProtocolError):
        b(m3)
    assert b._store["x"] == psk  # an AEAD failure burns nothing


def test_low_order_x25519_rejected():
    # non-contributory public keys are rejected at the DH exchange
    sk = _x25519_private_key(secrets.token_bytes(32))
    with pytest.raises(ValueError):
        sk.exchange(X25519PublicKey.from_public_bytes(b"\x00" * 32))


def test_foreign_mlkem_static_rejected_at_msg4(psk):
    # alice presents a static whose ML-KEM half she does not hold (foreign
    # mlkem half, her own x25519 half): she cannot decapsulate bob's ct_se
    # challenge, so bob's confirm tag does not match her root and she burns
    a, b = make_sessions(psk)
    # alice presents a static whose ML-KEM half she does not hold
    a._exchange._static_m = _public_from_idroot(secrets.token_bytes(32))[0]
    (m1,) = a()
    (m2,) = b(m1)
    (m3,) = a(m2)
    (vb,) = b(m3)
    (m4,) = b(vb)
    with pytest.raises(ProtocolError, match="responder tag"):
        a(m4)
    assert a._store == {}  # burned


def test_foreign_mlkem_static_rejected_responder(psk):
    # the symmetric check on the responder: bob cannot answer alice's
    # ct_es challenge for the static he presented -- his decapsulation
    # key does not match it, the inner box keyed with es does not open
    # for him, and he burns
    a, b = make_sessions(psk)
    (m1,) = a()
    (m2,) = b(m1)
    b._exchange._mlkem = MLKEM768PrivateKey.generate()  # not the presented key
    (m3,) = a(m2)
    with pytest.raises(ProtocolError, match="burned"):
        b(m3)
    assert b._store == {}  # burned


def test_credential_id(psk):
    # 16 bytes, recomputable from the stored record
    assert len(credential_id(psk)) == 16
    assert credential_id(psk) == credential_id(bytes(psk))
    a, b = make_sessions(psk)
    run_handshake(a, b, answer=None)
    blob = a.credential
    # stable across the verified flag (metadata may change)
    flagged = _CREDENTIAL.pack(*_CREDENTIAL.unpack(blob)[:5], True)
    assert credential_id(blob) == credential_id(flagged)
    # tied to the credential: a different record hashes differently
    assert credential_id(blob) != credential_id(psk)  # rotation rotates it
    assert credential_id(blob) != credential_id(b.credential)  # per-record
    forged = _CREDENTIAL.pack(secrets.token_bytes(32), *_CREDENTIAL.unpack(blob)[1:])
    assert credential_id(forged) != credential_id(blob)


def test_confirm_guards(psk):
    a, b = make_sessions(psk)
    with pytest.raises(ProtocolError, match="nothing to confirm"):
        a.confirm(True)  # handshake not complete
    run_handshake(a, b)  # verified by the registration answer already
    with pytest.raises(ProtocolError, match="nothing to confirm"):
        a.confirm(False)
    with pytest.raises(ProtocolError, match="nothing to verify"):
        _ = a.verification  # verified: no comparison value anymore