(* LO-KEX: Session Establishment (Theorems 1, 2b)
 *
 * Protocol (§4):
 *   Bob publishes bundle with σ_SPK = Sign(sk_IK_B, label ‖ pk_SPK_B)
 *   Alice verifies bundle, encapsulates to IK_B/SPK_B/OPK_B
 *   Alice signs SI: σ_SI = Sign(sk_IK_A, label ‖ SI)
 *   Bob verifies σ_SI, decapsulates, derives same (rk, ek)
 *
 * Simplifications:
 *   - HybridSig → single EUF-CMA scheme (§2.2: "may be treated as single")
 *   - Three X-Wing KEMs → three independent IND-CCA2 KEMs
 *   - HKDF → PRF keyed by ss_IK with (ss_SPK, ss_OPK, pks) as input
 *   - First-message AEAD omitted (Theorem 1 is about session key, not message)
 *   - σ_SPK signed by Bob's signing key (same as IK in the hybrid scheme)
 *   - Alice uses a SEPARATE signing key (models the distinct sk_IK_A[Ed25519+ML-DSA])
 *)

proof {
  crypto uf_cma(sigA_sign);
  simplify;
  success
}

param N_A.
param N_B.

(* ---------- Types ---------- *)

type kem_keyseed [large, fixed].
type kem_pkey [bounded].
type kem_skey [bounded].
type kem_secret [large, fixed].
type kem_ciphertext [bounded].
type kem_encapoutput [bounded].

type spk_keyseed [large, fixed].
type spk_pkey [bounded].
type spk_skey [bounded].
type spk_secret [large, fixed].
type spk_ciphertext [bounded].
type spk_encapoutput [bounded].

type opk_keyseed [large, fixed].
type opk_pkey [bounded].
type opk_skey [bounded].
type opk_secret [large, fixed].
type opk_ciphertext [bounded].
type opk_encapoutput [bounded].

type sig_keyseed [large, fixed].
type sig_pkey [bounded].
type sig_skey [bounded].
type sig_signature [bounded].

type sessionkey [large, fixed].

(* ---------- IK KEM (IND-CCA2) ---------- *)

proba P_kem_ik.
proba P_kem_ik_keycoll.
proba P_kem_ik_ctxtcoll.

expand IND_CCA2_KEM(
  kem_keyseed, kem_pkey, kem_skey,
  kem_secret, kem_ciphertext, kem_encapoutput,
  ik_pkgen, ik_skgen, ik_encap, ik_pair, ik_decap,
  injbot_ik, P_kem_ik, P_kem_ik_keycoll, P_kem_ik_ctxtcoll
).

(* ---------- SPK KEM (IND-CCA2) ---------- *)

proba P_kem_spk.
proba P_kem_spk_keycoll.
proba P_kem_spk_ctxtcoll.

expand IND_CCA2_KEM(
  spk_keyseed, spk_pkey, spk_skey,
  spk_secret, spk_ciphertext, spk_encapoutput,
  spk_pkgen, spk_skgen, spk_encap, spk_pair, spk_decap,
  injbot_spk, P_kem_spk, P_kem_spk_keycoll, P_kem_spk_ctxtcoll
).

(* ---------- OPK KEM (IND-CCA2) ---------- *)

proba P_kem_opk.
proba P_kem_opk_keycoll.
proba P_kem_opk_ctxtcoll.

expand IND_CCA2_KEM(
  opk_keyseed, opk_pkey, opk_skey,
  opk_secret, opk_ciphertext, opk_encapoutput,
  opk_pkgen, opk_skgen, opk_encap, opk_pair, opk_decap,
  injbot_opk, P_kem_opk, P_kem_opk_keycoll, P_kem_opk_ctxtcoll
).

(* ---------- Bob's signature (EUF-CMA) — signs SPK bundle ---------- *)

proba P_sig_B.
proba P_sig_B_keycoll.

expand UF_CMA_proba_signature(
  sig_keyseed, sig_pkey, sig_skey, bitstring, sig_signature,
  sigB_skgen, sigB_pkgen, sigB_sign, sigB_verify,
  P_sig_B, P_sig_B_keycoll
).

(* ---------- Alice's signature (EUF-CMA) — signs SessionInit ---------- *)

type sigA_keyseed [large, fixed].
type sigA_pkey [bounded].
type sigA_skey [bounded].
type sigA_signature [bounded].

proba P_sig_A.
proba P_sig_A_keycoll.

expand UF_CMA_proba_signature(
  sigA_keyseed, sigA_pkey, sigA_skey, bitstring, sigA_signature,
  sigA_skgen, sigA_pkgen, sigA_sign, sigA_verify,
  P_sig_A, P_sig_A_keycoll
).

(* ---------- HKDF as PRF ---------- *)

proba P_prf.

expand PRF_large(kem_secret, bitstring, sessionkey, kdf_kex, P_prf).

(* ---------- Domain separation constants ---------- *)

const lo_spk_sig_label: bitstring.
const lo_kex_init_sig_label: bitstring.

(* ---------- Events ---------- *)

(* Events bound on the signed session init content — what the signature
 * directly authenticates. rk is derived from this, but the correspondence
 * is over the authenticated payload. *)
event Alice_Init(sigA_pkey, sig_pkey, kem_ciphertext, spk_ciphertext, opk_ciphertext).
event Bob_Accept(sigA_pkey, sig_pkey, kem_ciphertext, spk_ciphertext, opk_ciphertext).

(* ---------- Security queries ---------- *)

(* Theorem 2b: Initiator authentication — if Bob accepts SI, Alice signed it *)
query pkA: sigA_pkey, pkB: sig_pkey,
      cik: kem_ciphertext, cspk: spk_ciphertext, copk: opk_ciphertext;
  event(Bob_Accept(pkA, pkB, cik, cspk, copk))
    ==> event(Alice_Init(pkA, pkB, cik, cspk, copk)).

(* Note: Injective variant not claimed — session-init replay is application-layer
 * (§7.5 A4). Alice may reuse the same bundle, producing identical SI contents.
 * Replay prevention is a caller obligation, not a cryptographic guarantee. *)

(* ---------- Channels ---------- *)

channel c_start, c_pub, c_alice_start, c_alice_out, c_bob_in.

(* ---------- Alice (Initiator) ---------- *)

let Alice(sk_sig_A: sigA_skey, pk_sig_A: sigA_pkey,
          pk_sig_B: sig_pkey, pk_ik_B: kem_pkey,
          pk_spk_B: spk_pkey, pk_opk_B: opk_pkey,
          sig_spk: sig_signature) =
  foreach i_a <= N_A do
  in(c_alice_start, ());
  (* §4.2: Verify Bob's SPK signature *)
  if sigB_verify((lo_spk_sig_label, pk_spk_B), pk_sig_B, sig_spk) = true then (
    (* §4.3 Step 2: Encapsulate to IK, SPK, OPK *)
    let ik_pair(ss_ik: kem_secret, c_ik: kem_ciphertext) = ik_encap(pk_ik_B) in
    let spk_pair(ss_spk: spk_secret, c_spk: spk_ciphertext) = spk_encap(pk_spk_B) in
    let opk_pair(ss_opk: opk_secret, c_opk: opk_ciphertext) = opk_encap(pk_opk_B) in
    (* §4.3 Step 3: Derive session key *)
    let rk: sessionkey = kdf_kex(ss_ik, (ss_spk, ss_opk, pk_sig_A, pk_sig_B)) in
    (* §4.3 Step 5: Sign session init *)
    let si: bitstring = (pk_sig_A, pk_sig_B, c_ik, c_spk, c_opk) in
    let sig_si: sigA_signature = sigA_sign((lo_kex_init_sig_label, si), sk_sig_A) in
    event Alice_Init(pk_sig_A, pk_sig_B, c_ik, c_spk, c_opk);
    out(c_alice_out, (si, sig_si, c_ik, c_spk, c_opk))
  ).

(* ---------- Bob (Responder) ---------- *)

let Bob(sk_ik_B: kem_skey, sk_spk_B: spk_skey, sk_opk_B: opk_skey,
        pk_sig_A: sigA_pkey, pk_sig_B: sig_pkey) =
  foreach i_b <= N_B do
  in(c_bob_in, (sig_si: sigA_signature,
                c_ik: kem_ciphertext, c_spk: spk_ciphertext,
                c_opk: opk_ciphertext));
  (* Bob reconstructs SI from the received components *)
  let si: bitstring = (pk_sig_A, pk_sig_B, c_ik, c_spk, c_opk) in
  (* §4.4 Step 2: Verify Alice's signature *)
  if sigA_verify((lo_kex_init_sig_label, si), pk_sig_A, sig_si) = true then (
    (* §4.4 Step 5: Decapsulate *)
    let injbot_ik(ss_ik: kem_secret) = ik_decap(c_ik, sk_ik_B) in
    let injbot_spk(ss_spk: spk_secret) = spk_decap(c_spk, sk_spk_B) in
    let injbot_opk(ss_opk: opk_secret) = opk_decap(c_opk, sk_opk_B) in
    (* §4.4 Step 7: Derive session key *)
    let rk: sessionkey = kdf_kex(ss_ik, (ss_spk, ss_opk, pk_sig_A, pk_sig_B)) in
    event Bob_Accept(pk_sig_A, pk_sig_B, c_ik, c_spk, c_opk)
  ).

(* ---------- Main process ---------- *)

process
  in(c_start, ());
  (* Bob: IK KEM key *)
  new ik_seed: kem_keyseed;
  let pk_ik_B = ik_pkgen(ik_seed) in
  let sk_ik_B = ik_skgen(ik_seed) in
  (* Bob: SPK *)
  new spk_seed: spk_keyseed;
  let pk_spk_B = spk_pkgen(spk_seed) in
  let sk_spk_B = spk_skgen(spk_seed) in
  (* Bob: OPK *)
  new opk_seed: opk_keyseed;
  let pk_opk_B = opk_pkgen(opk_seed) in
  let sk_opk_B = opk_skgen(opk_seed) in
  (* Bob: signing key (signs SPK bundle) *)
  new sigB_seed: sig_keyseed;
  let pk_sig_B = sigB_pkgen(sigB_seed) in
  let sk_sig_B = sigB_skgen(sigB_seed) in
  (* Alice: signing key (signs SessionInit) *)
  new sigA_seed: sigA_keyseed;
  let pk_sig_A = sigA_pkgen(sigA_seed) in
  let sk_sig_A = sigA_skgen(sigA_seed) in
  (* Bob signs SPK *)
  let sig_spk: sig_signature = sigB_sign((lo_spk_sig_label, pk_spk_B), sk_sig_B) in
  (* Publish everything *)
  out(c_pub, (pk_ik_B, pk_spk_B, pk_opk_B, pk_sig_B, pk_sig_A, sig_spk));
  (Alice(sk_sig_A, pk_sig_A, pk_sig_B, pk_ik_B, pk_spk_B, pk_opk_B, sig_spk)
   | Bob(sk_ik_B, sk_spk_B, sk_opk_B, pk_sig_A, pk_sig_B))