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

-- 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

axiom publicKey : KeyMode → Agent → Key
axiom injective_publicKey : ∀ {b c : KeyMode} {A A' : Agent},
  publicKey b A = publicKey c A' → b = c ∧ A = A'

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)

-- Axioms for private and public keys
axiom privateKey_neq_publicKey {b c : KeyMode} {A A' : Agent} :
    privateKey b A ≠ publicKey c A'

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
omit [InvKey] in
lemma publicKey_inject {b c : KeyMode} {A A' : Agent} :
    (publicKey b A = publicKey c A') ↔ (b = c ∧ A = A') := by
  grind[injective_publicKey]

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]; grind[privateKey_neq_publicKey];

lemma syKey_neq_priEK :
  K ∈ symKeys → K ≠ priEK A := by
  intro _ _
  have _ := not_symKeys_pubK (b := KeyMode.Encryption) (A := A)
  simp_all[symKeys, invKey_eq]

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

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

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

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

omit [InvKey] in
@[simp]
lemma publicKey_image_eq :
    (publicKey b x ∈ publicKey c '' AA) ↔ (b = c ∧ x ∈ AA) := by
  simp [Set.mem_image, publicKey_inject, 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]
 
-- 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.
axiom shrK : Agent → Key

axiom inj_shrK : Function.Injective shrK

-- All shared keys are symmetric
axiom sym_shrK : ∀ {A : Agent}, shrK A ∈ symKeys 

-- 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


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

omit [InvKey] in
@[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
@[simp]
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]
      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] <;>
  repeat rw[Set.singleton_def, parts_insert_Key, parts_empty] <;>
  simp

lemma Crypt_notin_initState {B : Agent} :
  Msg.Crypt K X ∉ parts ( initState B ) := by
  cases B <;> simp[initState, priEK, priSK, shrK] <;>
  apply And.intro <;> try apply And.intro
  all_goals repeat rw[Set.singleton_def, parts_insert_Key, parts_empty] <;>
  simp

@[simp]
lemma Crypt_notin_used_empty :
  Msg.Crypt K X ∉ used [] := by
  simp[used]; intro A; cases A <;> simp <;> apply And.intro <;> try apply And.intro
  all_goals (rw[Set.singleton_def, parts_insert_Key, parts_empty] ; simp)

-- 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 <;>
  try simp
  · left; left; right; exists Agent.Spy; apply And.intro; exact Spy_in_bad; rfl
  · left; left; left; exists Agent.Spy; apply And.intro; exact Spy_in_bad; rfl

@[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]
           cases b
           · right; exists A
           · left; right; exists A
  | 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
lemma Spy_spies_bad_privateKey { h : A ∈ bad } : Msg.Key (privateKey b A) ∈ spies evs := by
  induction evs with
  | nil => simp [spies, knows, pubSK, pubEK]; left; left
           cases b
           · right; exists A
           · left; exists A
  | 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]; 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

-- 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 Nonce_notin_initState {B : Agent} : Msg.Nonce N ∉ parts (initState B) := by
  cases B <;>
  simp [initState] <;> apply And.intro <;> try (apply And.intro)
  all_goals (rw[Set.singleton_def, parts_insert_Key, parts_empty]; simp)

@[simp]
lemma Nonce_notin_used_empty : Msg.Nonce N ∉ used [] := by
  simp [used]; intro A; cases A <;> simp <;>
  apply And.intro <;> try (apply And.intro)
  all_goals (rw[Set.singleton_def, parts_insert_Key, parts_empty]; 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)⟩

-- TODO is this really needed?
-- lemma Nonce_supply : Msg.Nonce (Classical.some (Nonce_supply_lemma.some_spec)) ∉ used evs := by
--   obtain ⟨N, h⟩ := Nonce_supply_lemma
--   exact h (Classical.some (Nonce_supply_lemma.some_spec)) (le_refl _)

-- 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] 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] 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