import InductiveVerification.Public

-- The Blanchet Protocol
namespace Blanchet

variable [InvKey]
variable [Bad]
variable [AgentKeys]
open Msg
open Event
open Bad
open HasInitState
open InvKey

-- Define the inductive set `blanchet`
inductive blanchet : List Event → Prop
| Nil : blanchet []
| Fake : blanchet evsf →
    X ∈ synth (analz (spies evsf)) →
    blanchet (Says Agent.Spy B X :: evsf)
| B1 : blanchet evs1 →
    sk ∈ symKeys →
    -- sk is fresh
    Key sk ∉ used evs1 ∪ Key '' keysFor (used evs1) →
    blanchet (Says A B (Crypt (pubEK B) (Sign (priSK A) ⦃Agent B, Key sk⦄)) :: evs1)
| B2 : blanchet evs2 →
    Nonce s ∉ used evs2 →
    Says A B (Crypt (pubEK B) (Sign (priSK A) ⦃Agent B, Key k⦄)) ∈ evs2 →
    blanchet (Says B A (Crypt k (Nonce s)) :: evs2)

-- A "possibility property": there are traces that reach the end
theorem possibility_property :
  ∃ sk ∈ symKeys, ∃ s, ∃ evs, blanchet evs ∧ Says B A (Crypt sk s) ∈ evs
:= by
  obtain ⟨sk, _, _⟩ := symK_supply (evs := [])
  exists sk
  simp_all
  exists Nonce 0
  exists [
           Says B A (Crypt sk (Nonce 0)),
           Says A B (Crypt (pubEK B) (Sign (priSK A) ⦃Agent B, Key sk⦄)),
         ]
  constructor
  · apply blanchet.B2
    · apply_rules [ blanchet.B1, blanchet.Nil ]; simp_all
    · simp [used]
    · tauto
  · simp

-- Spy never sees another agent's private key unless it's bad at the start
@[simp, grind =]
theorem Spy_see_priEK {h : blanchet evs} :
  (Key (priEK A) ∈ parts (spies evs)) ↔ A ∈ bad := by
  constructor
  · induction h with
    | Nil => simp [ priEK, initState ]
    | Fake _ h =>
      intro h₁; simp at h₁
      apply Or.imp_left (f := Fake_parts_sing (h := h)) at h₁
      simp_all
    | B1 => simp_all; 
    | B2 => simp_all;
  · intro _; apply_rules [ parts_increasing, Spy_spies_bad_privateKey ]

@[simp]
theorem Spy_analz_priEK {h : blanchet evs} :
  Key (priEK A) ∈ analz (spies evs) ↔ A ∈ bad
:= by grind

@[simp, grind =]
theorem Spy_see_priSK {h : blanchet evs} :
  Key (priSK A) ∈ parts (spies evs) ↔ A ∈ bad
:= by
  constructor
  · induction h with
    | Nil => simp [ priSK, initState ]
    | Fake _ h =>
      intro h₁; simp at h₁
      apply Or.imp_left (f := Fake_parts_sing (h := h)) at h₁
      simp_all
    | B1 => simp_all; 
    | B2 => simp_all;
  · intro _; apply_rules [ parts_increasing, Spy_spies_bad_privateKey ]
  
@[simp]
theorem Spy_analz_priSK {h : blanchet evs} :
  Key (priSK A) ∈ analz (spies evs) ↔ A ∈ bad
:= by grind

-- Unicity for NS1: nonce NA identifies agents A and B
@[grind! .]
theorem unique_B1 { h : blanchet evs } :
  (Crypt (pubEK B) (Sign (priSK A) ⦃Agent B, Key sk⦄) ∈ parts (spies evs) →
  (Crypt (pubEK B') (Sign (priSK A') ⦃Agent B', Key sk⦄) ∈ parts (spies evs) →
  Key sk ∉ analz (spies evs) →
  A = A' ∧ B = B')) := by
  intro h₁ h₂ h₃
  induction h with
  | Nil => simp_all [ initState ]
  | Fake _ a a_ih =>
    apply mt (h₁ := analz_spies_mono) at h₃;
    simp [*, priSK] at *
    apply Or.imp_left (f := Fake_parts_sing (h := a)) at h₁
    apply Or.imp_left (f := Fake_parts_sing (h := a)) at h₂
    simp_all;
  | B1 =>
    simp [*] at *; expand_parts_element at h₁; expand_parts_element at h₂; grind
  | B2 => simp_all; grind

lemma keysFor_used_knows_Spy_helper :
  K ∈ keysFor (analz (spies evs)) → K ∈ keysFor (used evs)
:= by
  aapply keysFor_mono; apply analz_knows_Spy_subset_used

-- Spy does not see the key sent in B1 if A and B are secure
@[grind! .]
theorem Spy_not_see_sk { h : blanchet evs  }
  { not_bad_A : A ∉ bad }
  { not_bad_B : B ∉ bad } :
  Says A B (Crypt (pubEK B) (Sign (priSK A) ⦃Agent B, Key sk⦄)) ∈ evs →
  Key sk ∉ analz (spies evs) := by
  intro h₁ h₂
  induction h with
  | Nil => simp_all
  | Fake _ a => apply Fake_analz_insert at a; apply a at h₂; simp_all
  | B1 _ _ a =>
    simp_all; obtain ⟨a, b⟩ := a; rcases h₁ with (_ | h)
    · simp_all; apply a; aapply analz_knows_Spy_subset_used
    · have _ := h; apply Says_imp_used at h; apply used_parts_subset_parts at h;
      apply mt (h₁ := keysFor_used_knows_Spy_helper) at b
      apply analz_insert_Crypt_subset at h₂; simp at h₂
      apply analz_insert_Crypt_subset at h₂; simp at h₂
      rw [analz_insert_Key] at h₂
      cases h₂; all_goals simp_all[Set.subset_def]
  | B2 => apply analz_insert_Crypt_subset at h₂; simp at h₂; grind

-- Authentication for `A`: if she receives message 2 and has used `sk` to start
-- a run, then `B` has sent message 2.
theorem A_trusts_B2 {h : blanchet evs }
  { not_bad_A : A ∉ bad }
  { not_bad_B : B ∉ bad } :
  Says A B (Crypt (pubEK B) (Sign (priSK A) ⦃Agent B, Key sk⦄)) ∈ evs →
  Says B' A (Crypt sk (Number s)) ∈ evs →
  Says B A (Crypt sk (Number s)) ∈ evs
:= by
   intro h₁ h₂;
   apply Says_imp_parts_knows_Spy at h₂
   -- use unique_NA to show that B' = B
   induction h with
  | Nil => simp_all
  | Fake _ a =>
    have snssk := h₁; apply Spy_not_see_sk at snssk <;> try assumption
    apply mt (h₁ := analz_spies_mono) at snssk
    simp [*] at *
    cases h₁
    · simp_all
    · apply Or.imp_left (f := Fake_parts_sing (h := a)) at h₂; simp_all
    · aapply blanchet.Fake
  | B1 => simp [*] at *; expand_parts_element at h₂; simp_all [keysFor]; grind
  | B2 => simp [*] at *; grind
  
theorem sk_symmetric_B1 {h :blanchet evs }
  { not_bad_A : A ∉ bad } :
  Says A B (Crypt (pubEK B) (Sign (priSK A) ⦃Agent B, Key sk⦄)) ∈ evs →
  sk ∈ symKeys
:= by
  intro h₁; induction h <;> simp_all
  · grind
  · cases h₁ <;> simp_all

-- If the encrypted message appears then it originated with Alice in `NS1`
lemma B_trusts_B1 { h : blanchet evs}
  { not_bad_A : A ∉ bad } :
  Crypt (pubEK B) (Sign (priSK A) ⦃Agent B, Key sk⦄) ∈ parts (spies evs) →
  Key sk ∉ analz (spies evs) →
  Says A B (Crypt (pubEK B) (Sign (priSK A) ⦃Agent B, Key sk⦄)) ∈ evs
:= by
  intro h₁ h₂
  induction h with
  | Nil => simp [initState] at h₁
  | Fake _ a =>
    apply mt (h₁ := analz_spies_mono) at h₂
    simp at h₁; apply Or.imp_left (f := Fake_parts_sing (h := a)) at h₁; simp_all
  | B1 => apply mt (h₁ := analz_spies_mono) at h₂; simp_all; grind
  | B2 => apply mt (h₁ := analz_spies_mono) at h₂; simp_all;

-- Authenticity Properties obtained from `NS2`

-- B only respods to requests from A
theorem B_is_response { h : blanchet evs }
  { not_bad_B : B ∉ bad } :
  Says B A (Crypt sk (Nonce s)) ∈ evs →
  Says A B (Crypt (pubEK B) (Sign (priSK A) ⦃Agent B, Key sk⦄)) ∈ evs
:= by
  intro h₁
  induction h <;> simp_all <;> grind

-- `s` remains secret
theorem Spy_not_see_s { h : blanchet evs }
  { not_bad_A : A ∉ bad }
  { not_bad_B : B ∉ bad } :
  Says B A (Crypt sk (Nonce s)) ∈ evs →
  Nonce s ∉ analz (spies evs)
:= by
  intro h₁
  induction h with
  | Nil => simp_all
  | Fake _ a => intro h₂; apply Fake_analz_insert at a; apply a at h₂; simp_all;
  | B1 _ _ a =>
    simp at a; obtain ⟨a, b⟩ := a;
    simp [*, analz_insert_Crypt_element] at *
    have _ := h₁; apply Says_imp_parts_knows_Spy at h₁
    expand_parts_element at h₁; simp_all; intro h₁ h₂
    rw[analz_insert_Key] at h₂; simp at h₂; grind; 
    aapply mt (h₁ := keysFor_used_knows_Spy_helper)
  | B2 h not_used a a_ih =>
    simp at h₁; rcases h₁ with (_ | h) <;> simp_all[analz_insert_Crypt_element]
    . apply And.intro;
      · have b := a; apply Spy_not_see_sk at a; apply sk_symmetric_B1 at b
        all_goals simp_all
      · grind
    · apply Says_imp_used at h
      apply used_parts_subset_parts at h; rw[Set.subset_def] at h
      intro _ _; simp_all

end Blanchet