KEMAnonymity

abbrev KEM (PK SK C K : Type) := KEMScheme ProbComp K PK SK C

theorem probOutput_true_eq_one_of_imp {α : Type} (oa : ProbComp α) {p q : α → Bool}
    (h : ∀ a ∈ support oa, p a = true → q a = true)
    (hp : Pr[= true | (do let a ← oa; pure (p a))] = 1) :
    Pr[= true | (do let a ← oa; pure (q a))] = 1 := by

def run : ProbComp ((PK × SK) × (C × K) × Option K) := do
  let ks ← kem.keygen
  let ck ← kem.encaps ks.1
  let k? ← kem.decaps ks.2 ck.1
  pure (ks, ck, k?)

theorem correctExp_eq [DecidableEq K] :
    kem.CorrectExp = (do let a ← kem.run; pure (decide (a.2.2 = some a.2.1.2))) := by

def anonSetup : ProbComp (PK × PK) := do
  let (pk0, _) ← kem.keygen
  let (pk1, _) ← kem.keygen
  pure (pk0, pk1)

def anonBranch (b : Bool) (a : PK × PK) : ProbComp Bool := do
  let (c, _) ← kem.encaps (if b then a.2 else a.1)
  adv a.1 a.2 c

def AnonExp : ProbComp Bool := do
  let a ← kem.anonSetup
  let b ← ($ᵗ Bool)
  let z ← kem.anonBranch adv b a
  pure (b == z)

noncomputable def anonAdvantage : ℝ := (kem.AnonExp adv).boolBiasAdvantage

theorem anonAdvantage_eq_branchDistAdvantage :
    kem.anonAdvantage adv =
      (kem.anonSetup >>= kem.anonBranch adv true).boolDistAdvantage
      (kem.anonSetup >>= kem.anonBranch adv false) :=
  ProbComp.boolBiasAdvantage_bind_uniformBool_branch kem.anonSetup (kem.anonBranch adv)

def StealthScheme.ofKEM (kem : KEM PK SK C K) : StealthScheme PK SK C where
  keygen := kem.keygen
  announce pk := do
    let (c, _) ← kem.encaps pk
    pure c
  scan sk c := do
    let k ← kem.decaps sk c
    pure k.isSome

theorem unlinkAdvantage_ofKEM_eq_anonAdvantage
    (kem : KEM PK SK C K) (adv : StealthScheme.UnlinkAdv PK C) :
    (StealthScheme.ofKEM kem).unlinkAdvantage adv = kem.anonAdvantage adv := by

def StealthScheme.ofKEMFull [DecidableEq Aux] (kem : KEM PK SK C K)
    (auxGen : K → PK → Aux) : StealthScheme PK (SK × PK) (C × Aux) where
  keygen := do
    let ks ← kem.keygen
    pure (ks.1, (ks.2, ks.1))
  announce pk := do
    let ck ← kem.encaps pk
    pure (ck.1, auxGen ck.2 pk)
  scan sk ca := do
    let k? ← kem.decaps sk.1 ca.1
    pure (k?.elim false fun k => decide (ca.2 = auxGen k sk.2))

theorem perfectlyComplete_ofKEMFull [DecidableEq K] [DecidableEq Aux]
    (kem : KEM PK SK C K) (auxGen : K → PK → Aux)
    (hkem : kem.PerfectlyCorrect ProbCompRuntime.probComp) :
    (StealthScheme.ofKEMFull kem auxGen).PerfectlyComplete := by

theorem unlinkSetup_ofKEMFull [DecidableEq Aux] (kem : KEM PK SK C K)
    (auxGen : K → PK → Aux) :
    (StealthScheme.ofKEMFull kem auxGen).unlinkSetup = kem.anonSetup := by

def StealthScheme.UnlinkAdv.cipherOf
    (adv : StealthScheme.UnlinkAdv PK (C × Aux)) (auxGen : K → PK → Aux) :
    StealthScheme.UnlinkAdv PK C :=
  fun pk0 pk1 c => do
    let k' ← ($ᵗ K)
    adv pk0 pk1 (c, auxGen k' pk0)

def randAuxBranch (b : Bool) (a : PK × PK) : ProbComp Bool := do
  let (c, _) ← kem.encaps (if b then a.2 else a.1)
  let k' ← ($ᵗ K)
  adv a.1 a.2 (c, auxGen k' (if b then a.2 else a.1))

theorem randAuxBranch_false :
    randAuxBranch kem auxGen adv false = kem.anonBranch (adv.cipherOf auxGen) false := rfl

noncomputable def sharedSecretHiding (b : Bool) : ℝ :=
  ((StealthScheme.ofKEMFull kem auxGen).unlinkSetup >>=
      (StealthScheme.ofKEMFull kem auxGen).unlinkBranch adv b).boolDistAdvantage
    ((StealthScheme.ofKEMFull kem auxGen).unlinkSetup >>=
      randAuxBranch kem auxGen adv b)

noncomputable def auxKeyIndependence : ℝ :=
  ((StealthScheme.ofKEMFull kem auxGen).unlinkSetup >>=
      randAuxBranch kem auxGen adv true).boolDistAdvantage
    (kem.anonSetup >>= kem.anonBranch (adv.cipherOf auxGen) true)

theorem unlinkAdvantage_ofKEMFull_le :
    (StealthScheme.ofKEMFull kem auxGen).unlinkAdvantage adv ≤
      sharedSecretHiding kem auxGen adv true
      + auxKeyIndependence kem auxGen adv
      + kem.anonAdvantage (adv.cipherOf auxGen)
      + sharedSecretHiding kem auxGen adv false := by