import InductiveVerification.Event
import Init.Data.Nat.Lemmas
import Mathlib.Data.Nat.Basic
import Mathlib.Data.Nat.Dist
import Mathlib.Order.Defs.PartialOrder

-- Theory of Public Keys (common to all public-key protocols)

-- Private and public keys; initial states of agents

variable [InvKey]
open InvKey

lemma invKey_K {K : Key} : K ∈ symKeys → invKey K = K := by
  intro h
  simp [symKeys] at h
  exact h

-- Asymmetric Keys
inductive KeyMode
  | Signature
  | Encryption

-- TODO replace axioms with classes
-- Also make sure that there are infinite key supplies
class AgentKeys where
  publicKey : KeyMode → Agent → Key
  injective_publicKey : ∀ {b c : KeyMode} {A A' : Agent},
    publicKey b A = publicKey c A' → b = c ∧ A = A'
  privateKey_neq_publicKey : invKey (publicKey b A) ≠ publicKey c A'
  not_surjective_publicKey_asymK : (¬all_symmetric) →
    (∀ K : Key, ∃ K' > K,
    K' ∉ symKeys ∧
    (∀ b : KeyMode, ∀ A : Agent,
    publicKey b A ≠ K' ∧ invKey (publicKey b A) ≠ K'))
-- Symmetric Keys
-- For some protocols, it is convenient to equip agents with symmetric as
-- well as asymmetric keys.  The theory ‹Shared› assumes that all keys
-- are symmetric.
  shrK : Agent → Key
  inj_shrK : Function.Injective shrK
  sym_shrK : ∀ {A : Agent}, shrK A ∈ symKeys 
  not_surjective_shrK_symK : ∀ K : Key, ∃ K' > K,
    K' ∈ symKeys ∧ (∀ A : Agent, shrK A ≠ K')

open AgentKeys
variable [AgentKeys]

noncomputable abbrev pubEK (A : Agent) : Key := publicKey KeyMode.Encryption A
noncomputable abbrev pubSK (A : Agent) : Key := publicKey KeyMode.Signature A
noncomputable abbrev privateKey (b : KeyMode) (A : Agent) : Key := invKey (publicKey b A)
noncomputable abbrev priEK (A : Agent) : Key := privateKey KeyMode.Encryption A
noncomputable abbrev priSK (A : Agent) : Key := privateKey KeyMode.Signature A
noncomputable abbrev pubK (A : Agent) : Key := pubEK A
noncomputable abbrev priK (A : Agent) : Key := invKey (pubEK A)

abbrev Sign (K : Key) (M : Msg) := ⦃Msg.Crypt K M, M⦄

attribute [simp] pubEK
attribute [simp] pubSK
-- attribute [simp] priEK
-- attribute [simp] priSK
attribute [simp] privateKey_neq_publicKey

-- Axioms for private and public keys
@[simp]
lemma publicKey_neq_privateKey {b c : KeyMode} {A A' : Agent} :
    publicKey b A ≠ privateKey c A' := by
  exact privateKey_neq_publicKey.symm

-- Basic properties of pubK and priK
@[simp]
lemma publicKey_inject {b c : KeyMode} {A A' : Agent} :
    (publicKey b A = publicKey c A') ↔ (b = c ∧ A = A') := by
  grind[injective_publicKey]

omit [AgentKeys] in
lemma invKey_injective: Function.Injective invKey := by
  intro _ _ _
  simp_all[invKey_eq]

lemma not_symKeys_pubK {b : KeyMode} {A : Agent} :
    publicKey b A ∉ symKeys :=
by
  grind[symKeys, privateKey_neq_publicKey]

lemma not_symKeys_priK {b : KeyMode} {A : Agent} :
    privateKey b A ∉ symKeys := by
  simp [symKeys, privateKey, invKey_eq, privateKey_neq_publicKey]
  
@[simp]
lemma pubK_neq_symK {b : KeyMode} {A : Agent} {h : K ∈ symKeys} :
   publicKey b A ≠ K
:= by
  intro h₁; have h₂ := not_symKeys_pubK (b := b) (A := A); simp_all

@[simp]
lemma priK_neq_symK {b : KeyMode} {A : Agent} {h : K ∈ symKeys} :
   privateKey b A ≠ K
:= by
  intro h₁; have h₂ := not_symKeys_priK (b := b) (A := A); simp_all

@[simp]
lemma symK_neq_pubK {b : KeyMode} {A : Agent} {h : K ∈ symKeys} :
   K ≠ publicKey b A
:= by intro h₁; aapply pubK_neq_symK; simp_all

@[simp]
lemma symK_neq_priK {b : KeyMode} {A : Agent} {h : K ∈ symKeys} :
   K ≠ privateKey b A
:= by intro h₁; aapply priK_neq_symK; simp_all

omit [AgentKeys] in
lemma symKeys_neq_imp_neq :
  ((K ∈ symKeys) ≠ (K' ∈ symKeys)) → K ≠ K' := by
  intro h eq
  rw[eq] at h
  contradiction

omit [AgentKeys] in
@[simp]
lemma symKeys_invKey_iff : (invKey K ∈ symKeys) = (K ∈ symKeys) := by
  simp [symKeys, invKey_eq]

omit [AgentKeys] in
lemma analz_symKeys_Decrypt :
    Msg.Crypt K X ∈ analz H → K ∈ symKeys → Msg.Key K ∈ analz H → X ∈ analz H := by
  simp [symKeys]
  intro _ _ _
  aapply analz.decrypt; simp_all

-- "Image" equations that hold for injective functions

omit [AgentKeys] in
@[simp]
lemma invKey_image_eq : (invKey x ∈ invKey '' A) ↔ (x ∈ A) := by
  simp [Set.mem_image]

@[simp]
lemma publicKey_image_eq :
    (publicKey b x ∈ publicKey c '' AA) ↔ (b = c ∧ x ∈ AA) := by
  simp [Set.mem_image, And.comm, Eq.comm]

@[simp]
lemma privateKey_notin_image_publicKey :
    privateKey b x ∉ publicKey c '' AA := by
  simp[publicKey_neq_privateKey]

@[simp]
lemma privateKey_image_eq :
    (privateKey b A ∈ invKey '' (publicKey c '' AS)) ↔ (b = c ∧ A ∈ AS) := by
  rw[←publicKey_image_eq]; simp [Set.mem_image, privateKey]

@[simp]
lemma publicKey_notin_image_privateKey :
    publicKey b A ∉ invKey '' ( publicKey c '' AS ) := by
  simp [privateKey_neq_publicKey]
 
-- Injectiveness: Agents' long-term keys are distinct.
@[simp]
lemma invKey_shrK :
    invKey (shrK A) = shrK A := by
  simp [invKey_K, sym_shrK]
  
lemma analz_shrK_Decrypt :
    Msg.Crypt (shrK A) X ∈ analz H → Msg.Key (shrK A) ∈ analz H → X ∈ analz H :=
by
  intro _ _
  aapply analz.decrypt; rw[invKey_shrK]; assumption
  
omit [AgentKeys] in
lemma analz_Decrypt' :
    Msg.Crypt K X ∈ analz H → K ∈ symKeys → Msg.Key K ∈ analz H → X ∈ analz H := by
  intro _ _ _
  aapply analz.decrypt; simp_all[symKeys]

@[simp]
lemma priK_neq_shrK {A : Agent} : shrK A ≠ privateKey b C := by
  intro h
  apply not_symKeys_priK
  rw[←h]
  exact sym_shrK

@[simp]
lemma shrK_neq_priK : privateKey b C ≠ shrK A := by
  exact priK_neq_shrK.symm

@[simp]
lemma pubK_neq_shrK : shrK A ≠ publicKey b C := by
  intro h
  apply not_symKeys_pubK
  rw[←h]
  exact sym_shrK

@[simp]
lemma shrK_neq_pubK : publicKey b C ≠ shrK A := by
  exact pubK_neq_shrK.symm

@[simp]
lemma priEK_noteq_shrK : priEK A ≠ shrK B := by
  simp

@[simp]
lemma publicKey_notin_image_shrK : publicKey b x ∉ shrK '' AA := by
  simp

@[simp]
lemma privateKey_notin_image_shrK : privateKey b x ∉ shrK '' AA := by
  simp

@[simp]
lemma shrK_notin_image_publicKey : shrK x ∉ publicKey b '' AA := by
  simp

@[simp]
lemma shrK_notin_image_privateKey :
  shrK x ∉ (invKey '' ((publicKey b) '' AA )) := by
  simp

@[simp]
lemma shrK_image_eq : (shrK x ∈ shrK '' AA) ↔ (x ∈ AA) := by
  grind[inj_shrK]

attribute [simp] invKey_K

variable [Bad]
open Bad
-- Fill in definition for Initial States of Agents
instance : HasInitState Agent where
  initState
    | Agent.Server =>
      {Msg.Key (priEK Agent.Server), Msg.Key (priSK Agent.Server)} ∪
      (Msg.Key '' Set.range pubEK) ∪ (Msg.Key '' Set.range pubSK) ∪ (Msg.Key '' Set.range shrK)
    | Agent.Friend i =>
      {Msg.Key (priEK (Agent.Friend i)), Msg.Key (priSK (Agent.Friend i)), Msg.Key (shrK (Agent.Friend i))} ∪
      (Msg.Key '' Set.range pubEK) ∪ (Msg.Key '' Set.range pubSK)
    | Agent.Spy =>
      (Msg.Key '' (invKey '' (pubEK '' bad))) ∪ (Msg.Key '' ( invKey '' ( pubSK '' bad  ))) ∪
      (Msg.Key '' ( shrK '' bad )) ∪
      (Msg.Key '' Set.range pubEK) ∪ (Msg.Key '' Set.range pubSK)

open HasInitState

-- These lemmas allow reasoning about `used evs` rather than `knows Spy evs`,
-- which is useful when there are private Notes. Because they depend upon the
-- definition of `initState`, they cannot be moved up.

lemma used_parts_subset_parts :
  X ∈ used evs → ( parts {X} ⊆ used evs ) := by
  induction evs with
  | nil => 
    simp[used]; intro A h₁ X h₂; simp; exists A
    cases A
    all_goals (
      simp_all[-parts_union, initState]
      apply_rules [parts_trans, h₂, Set.singleton_subset_iff.mpr]
    )
  | cons e evs ih =>
    intro h₁ X
    cases e <;> simp[used] <;> simp[used] at h₁ <;> try aapply ih;
    all_goals (
    cases h₁ with
      | inl h₁ => intro _; left; apply_rules [parts_trans, h₁, Set.singleton_subset_iff.mpr]
      | inr h₁ => intro h₂; right; apply ih at h₁; aapply h₁
    )

lemma MPair_used_D :
    ⦃X, Y⦄ ∈ used evs → (X ∈ used evs ∧ Y ∈ used evs) := by
  intro h
  apply used_parts_subset_parts (X := Msg.MPair X Y) (evs := evs) at h
  rw[Set.singleton_def, parts_insert_MPair] at h
  simp at h
  apply And.intro <;> apply h <;> tauto

lemma MPair_used {P : Prop} :
    ⦃X, Y⦄ ∈ used evs →
    (X ∈ used evs → Y ∈ used evs → P) →
    P := by
  intro hXY hP
  have ⟨hX, hY⟩ := MPair_used_D hXY
  exact hP hX hY

-- Define `@[simp]` lemmas for `initState` for each case of `Agent`
@[simp]
lemma keysFor_parts_initState {C : Agent} :
    keysFor (parts (initState C)) = ∅ := by
  cases C <;>
  simp[initState, keysFor]
  
@[simp]
lemma keysFor_used_empty :
  keysFor (used []) = ∅
:= by
  rw[used, keysFor_iunion]; simp;

lemma Crypt_notin_initState {B : Agent} :
  Msg.Crypt K X ∉ parts ( initState B ) := by
  cases B <;> simp[initState, priEK, priSK]

@[simp]
lemma Crypt_notin_used_empty :
  Msg.Crypt K X ∉ used [] := by
  simp[used]; intro A; cases A <;> simp[initState]

-- Basic properties of shrK

-- Agents see their own shared keys
-- iff
@[simp]
lemma shrK_in_initState {A : Agent} :
    Msg.Key (shrK A) ∈ initState A := by
  induction A <;> simp [HasInitState.initState, initState]
  exists Agent.Spy; simp[Spy_in_bad]

-- iff
@[simp]
lemma shrK_in_knows {A : Agent} : Msg.Key (shrK A) ∈ knows A evs := by
  apply initState_subset_knows
  exact shrK_in_initState

-- iff
@[simp]
lemma shrK_in_used {A : Agent} : Msg.Key (shrK A) ∈ used evs := by
  apply_rules [initState_into_used, parts.inj, shrK_in_initState]

-- Fresh keys never clash with long-term shared keys

-- Used in parts_induct_tac and analz_Fake_tac to distinguish session keys
-- from long-term shared keys
@[simp]
lemma Key_not_used {K : Key} : Msg.Key K ∉ used evs → K ∉ Set.range shrK := by
  simp
  intro h₁ _ h₂
  apply h₁
  rw[←h₂]
  apply shrK_in_used

lemma shrK_neq {K : Key} {B : Agent} : Msg.Key K ∉ used evs → shrK B ≠ K := by
  intro h h'
  apply h
  rw [←h']
  exact shrK_in_used

@[simp]
lemma neq_shrK {K : Key} {B : Agent} : Msg.Key K ∉ used evs → K ≠ shrK B := by
  intro h; apply shrK_neq at h; exact h.symm

-- Function spies

omit [InvKey] [Bad] in
lemma not_SignatureE {b : KeyMode} : b ≠ KeyMode.Signature → b = KeyMode.Encryption := by
  cases b <;> simp
  
-- Agents see their own private keys
@[simp]
lemma priK_in_initState {b : KeyMode} {A : Agent} :
    Msg.Key (privateKey b A) ∈ initState A := by
  induction A <;>
  simp [HasInitState.initState, initState, privateKey, pubEK, pubSK] <;>
  cases b <;> simp

@[simp]
lemma publicKey_in_initState {b : KeyMode} {A : Agent} {B : Agent} :
    Msg.Key (publicKey b A) ∈ initState B := by
  induction B <;>
  simp [HasInitState.initState, initState, priEK, priSK, pubEK, pubSK] <;>
  cases b <;>
  simp
  
-- All public keys are visible
@[simp]
lemma spies_pubK : Msg.Key (publicKey b A) ∈ spies evs := by
  induction evs with
  | nil => simp [spies, knows, initState]
           cases b <;> tauto
  | cons e evs ih =>
    cases e <;> rw [spies] <;> apply knows_subset_knows_Cons <;> assumption

@[simp]
lemma analz_spies_pubK : Msg.Key (publicKey b A) ∈ analz (spies evs) := by
  exact analz.inj spies_pubK

-- Spy sees private keys of bad agents
@[grind .]
lemma Spy_spies_bad_privateKey { h : A ∈ bad } : Msg.Key (privateKey b A) ∈ spies evs := by
  induction evs with
  | nil => simp_all [spies, knows, initState]; cases b <;> tauto
  | cons e evs ih =>
    cases e <;> rw[spies] <;> aapply knows_subset_knows_Cons

-- Spy sees long-term shared keys of bad agents
lemma Spy_spies_bad_shrK {h : A ∈ bad} : Msg.Key (shrK A) ∈ spies evs := by
  induction evs with
  | nil => simp [spies, knows, initState]; exists A
  | cons e evs ih =>
    cases e <;> rw [spies] <;> aapply knows_subset_knows_Cons

@[simp]
lemma publicKey_into_used : Msg.Key (publicKey b A) ∈ used evs := by
  aapply initState_into_used
  apply parts_increasing
  exact publicKey_in_initState

@[simp]
lemma privateKey_into_used : Msg.Key (privateKey b A) ∈ used evs := by
  aapply initState_into_used
  apply parts_increasing
  exact priK_in_initState
  
@[simp]
lemma shrK_into_used: Msg.Key (shrK A) ∈ used evs := by
  aapply initState_into_used
  apply parts_increasing
  exact shrK_in_initState

@[grind .]
lemma analz_priK_Decrypt :
  Msg.Crypt (priSK A) X ∈ analz (spies evs) → X ∈ analz (spies evs) :=
by
  intro h; aapply analz.decrypt
  simp[priSK, privateKey, invKey_spec]

-- For case analysis on whether or not an agent is compromised
lemma Crypt_Spy_analz_bad :
    Msg.Crypt (shrK A) X ∈ analz (knows Agent.Spy evs) → A ∈ bad → X ∈ analz (knows Agent.Spy evs) := by
    intro h₁ h₂
    aapply analz.decrypt
    apply analz_increasing
    simp[invKey_shrK]
    aapply Spy_spies_bad_shrK

@[simp]
lemma Crypt_synth_pubK :
  (Msg.Crypt (pubEK A) X ∈ synth (spies evs)) ↔
  (Msg.Crypt (pubEK A) X ∈ (spies evs) ∨ ( X ∈ synth (spies evs))) :=
by simp[Crypt_synth_EK]; 

@[simp]
lemma Crypt_synth_analz_pubK :
  (Msg.Crypt (pubEK A) X ∈ synth (analz (spies evs))) ↔
  (Msg.Crypt (pubEK A) X ∈ (analz (spies evs)) ∨ ( X ∈ synth (analz (spies evs)))) :=
by simp[Crypt_synth_EK]; 

@[simp]
lemma Crypt_synth_priK :
  (Msg.Crypt (priSK A) X ∈ synth (spies evs)) ↔
  (Msg.Crypt (priSK A) X ∈ spies evs ∨
  (Msg.Key (priSK A) ∈ spies evs ∧ X ∈ synth (spies evs))) :=
by simp[Crypt_synth_EK]

@[simp]
lemma Crypt_synth_analz_priK :
  (Msg.Crypt (priSK A) X ∈ synth (analz (spies evs))) ↔
  (Msg.Crypt (priSK A) X ∈ analz (spies evs) ∨
  (Msg.Key (priSK A) ∈ analz (spies evs) ∧ X ∈ synth (analz (spies evs)))) :=
by simp[Crypt_synth_EK]; 

@[grind .]
lemma Crypt_synth_analz_priK_decrypt :
  (Msg.Crypt (priSK A) X ∈ synth (analz (spies evs))) →
  (X ∈ analz (spies evs) ∨
  (Msg.Key (priSK A) ∈ analz (spies evs) ∧ X ∈ synth (analz (spies evs)))) :=
by simp[Crypt_synth_EK]; grind

@[simp]
lemma Sign_synth_priK :
  (Sign (priSK A) X ∈ synth (analz (spies evs))) ↔
  (Msg.Crypt (priSK A) X ∈ analz (spies evs) ∨ Msg.Key (priSK A) ∈ analz (spies evs)) ∧ 
  X ∈ synth (analz (spies evs))
:= by
  simp[Sign, Crypt_synth_EK]; grind

@[simp]
lemma Nonce_notin_initState {B : Agent} : Msg.Nonce N ∉ parts (initState B) := by
  cases B <;>
  simp [initState]

@[simp]
lemma Nonce_notin_used_empty : Msg.Nonce N ∉ used [] := by
  simp [used, initState]; intro A; cases A <;> simp

-- Supply fresh nonces for possibility theorems
lemma Nonce_supply_lemma : ∃ N, ∀ n, N ≤ n → Msg.Nonce n ∉ used evs := by
  induction evs with
  | nil =>
    use 0; simp
  | cons e evs ih =>
    obtain ⟨N₁, hN⟩ := ih
    cases e with
    | Says _ _ m => simp[used]
                    have ns := msg_Nonce_supply (msg := m)
                    obtain ⟨N₂, ns⟩ := ns
                    exists Nat.max N₁ N₂
                    simp_all
    | Notes _ m => simp[used]
                   have ns := msg_Nonce_supply (msg := m)
                   obtain ⟨N₂, ns⟩ := ns
                   exists Nat.max N₁ N₂
                   simp_all
    | Gets => simp[used]; exists N₁


lemma Nonce_supply1 : ∃ N, Msg.Nonce N ∉ used evs := by
  obtain ⟨N, h⟩ := Nonce_supply_lemma
  exact ⟨N, h N (le_refl N)⟩
  
lemma symK_supply_lemma : ∀ K, ∃ K' > K,
  K' ∈ symKeys ∧ Msg.Key K' ∉ used evs ∪ Msg.Key '' keysFor (used evs) :=
by
  induction evs with
  | nil =>
    intro K
    have exK := not_surjective_shrK_symK (K := K);
    rcases exK with ⟨K' , ⟨_, symK, exK⟩⟩;
    exists K';
    apply And.intro
    · assumption
    · apply And.intro
      · assumption
      · simp_all[used]; intro A; cases A <;>
        simp_all[initState, pubK_neq_symK, priK_neq_symK, symK_neq_priK]
        rw[Eq.comm]; apply exK
  | cons e evs ih =>
    intro K
    cases e with
    | Says _ _ m =>
      simp[used];
      obtain ⟨K₁, ks⟩ := msg_Key_supply (msg := m); 
      obtain ⟨K', _, _, _⟩ := ih ( K := Nat.max K₁ K); exists K'; simp_all; grind
    | Notes A m => 
      simp[used];
      obtain ⟨K₁, ks⟩ := msg_Key_supply (msg := m); 
      obtain ⟨K', _, _, _⟩ := ih ( K := Nat.max K₁ K); exists K';
      apply ks at K'; by_cases h : A ∈ bad <;> simp_all <;> grind
    | Gets => exact ih (K := K)
  
lemma symK_supply : ∃ K ∈ symKeys,
  Msg.Key K ∉ used evs ∪ Msg.Key '' keysFor (used evs) :=
by
  obtain ⟨K, _, _, _⟩ := symK_supply_lemma (K := 0) (evs := evs)
  exists K

-- Specialized Rewriting for Theorems About `analz` and Image
omit [InvKey] [Bad] in
lemma insert_Key_singleton : insert (Msg.Key K) H = Msg.Key '' {K} ∪ H := by
  simp

omit [InvKey] [Bad] in
@[simp]
lemma insert_Key_image : insert (Msg.Key K) (Msg.Key '' KK ∪ C) = Msg.Key '' (insert K KK) ∪ C := by
  rw[insert_Key_singleton, Set.image_insert_eq, Set.insert_eq, Set.union_assoc, Set.image_singleton]

omit [Bad] [AgentKeys] in
lemma Crypt_imp_keysFor :
    Msg.Crypt K X ∈ H → K ∈ symKeys → K ∈ keysFor H := by
  intro h₁ h₂
  apply invKey_K at h₂
  rw[←h₂]
  aapply Crypt_imp_invKey_keysFor

-- Lemma for the trivial direction of the if-and-only-if of the 
-- Session Key Compromise Theorem
omit [Bad] [AgentKeys] in
@[simp]
lemma analz_image_freshK_lemma :
    ((Msg.Key K ∈ analz (Msg.Key '' nE ∪ H)) →
    (K ∈ nE ∨ Msg.Key K ∈ analz H)) →
    (Msg.Key K ∈ analz (Msg.Key '' nE ∪ H)) = (K ∈ nE ∨ Msg.Key K ∈ analz H) :=
by
  intro h
  ext; constructor; assumption;
  intro h₁; simp_all; cases h₁
  · apply analz_increasing; left; simp; assumption
  · apply analz_union; right; assumption
