import InductiveVerification.Public

-- The Needham-Schroeder Public-Key Protocol
namespace NS_Public

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

attribute [-simp] Key.injEq

-- Define the inductive set `ns_public`
inductive ns_public : List Event → Prop
| Nil : ns_public []
| Fake : ns_public evsf →
    X ∈ synth (analz (spies evsf)) →
    ns_public (Says Agent.Spy B X :: evsf)
| NS1 : ns_public evs1 →
    Nonce NA ∉ used evs1 →
    ns_public (Says A B (Crypt (pubEK B) ⦃Nonce NA, Agent A⦄) :: evs1)
| NS2 : ns_public evs2 →
    Nonce NB ∉ used evs2 →
    Says A' B (Crypt (pubEK B) ⦃Nonce NA, Agent A⦄) ∈ evs2 →
    ns_public (Says B A (Crypt (pubEK A) ⦃Nonce NA, Nonce NB, Agent B⦄) :: evs2)
| NS3 : ns_public evs3 →
    Says A B (Crypt (pubEK B) ⦃Nonce NA, Agent A⦄) ∈ evs3 →
    Says B' A (Crypt (pubEK A) ⦃Nonce NA, Nonce NB, Agent B⦄) ∈ evs3 →
    ns_public (Says A B (Crypt (pubEK B) (Nonce NB)) :: evs3)

-- A "possibility property": there are traces that reach the end
theorem possibility_property :
  ∃ NB, ∃ evs, ns_public evs ∧ Says A B (Crypt (pubEK B) (Nonce NB)) ∈ evs := by
  exists 1
  exists [ Says A B (Crypt (pubEK B) (Nonce 1)),
           Says B A (Crypt (pubEK A) ⦃Nonce 0, Nonce 1, Agent B⦄),
           Says A B (Crypt (pubEK B) ⦃Nonce 0, Agent A⦄),
         ]
  constructor
  · apply ns_public.NS3
    · apply ns_public.NS2
      · apply_rules[ns_public.NS1, ns_public.Nil, Nonce_notin_used_empty]
      · simp
      · tauto
    all_goals tauto
  · simp

-- Spy never sees another agent's private key unless it's bad at the start
@[simp]
theorem Spy_see_priEK {h : ns_public evs} :
  (Key (priEK A) ∈ parts (spies evs)) ↔ A ∈ bad := by
  constructor
  · induction h with
    | Nil => simp[spies, knows, initState, pubEK, priEK, pubSK, Key.injEq]
    | Fake _ h =>
      apply Fake_parts_sing at h
      intro h₁; simp at h₁; apply Fake_parts_sing_helper (h := h) at h₁
      simp_all
    | NS1 => simp_all
    | NS2 => simp_all
    | NS3 => simp_all
  · intro _; apply_rules [ parts_increasing, Spy_spies_bad_privateKey ]

@[simp]
theorem Spy_analz_priEK {h : ns_public evs} :
  Key (priEK A) ∈ analz (spies evs) ↔ A ∈ bad := by
  constructor
  · intro h₁; apply analz_subset_parts at h₁; aapply Spy_see_priEK.mp
  · intro h₁; apply analz_increasing; aapply Spy_spies_bad_privateKey

-- It is impossible to re-use a nonce in both NS1 and NS2, provided the nonce is secret
theorem no_nonce_NS1_NS2 { evs: List Event} { h : ns_public evs  } :
  (Crypt (pubEK C) ⦃NA', Nonce NA, Agent D⦄ ∈ parts (spies evs) →
  (Crypt (pubEK B) ⦃Nonce NA, Agent A⦄ ∈ parts (spies evs) →
  Nonce NA ∈ analz (spies evs))) := by
  intro h₁ h₂
  induction h with
  | Nil => simp[spies, knows] at h₂
  | Fake _ h ih =>
    apply analz_spies_mono
    simp [*] at *
    apply Fake_parts_sing at h
    apply Fake_parts_sing_helper (h := h) at h₁
    apply Fake_parts_sing_helper (h := h) at h₂
    simp [*] at *; grind[analz_subset_parts]
  | NS1 _ nonce_not_used =>
    apply analz_spies_mono
    simp [*] at *
    expand_parts_element at h₁; expand_parts_element at h₂;
    grind [ parts_knows_Spy_subset_used ]
  | NS2 _ nonce_not_used =>
    apply analz_spies_mono
    simp [*] at *
    expand_parts_element at h₂;
    grind [ parts_knows_Spy_subset_used ]
  | NS3 _ _ _ a_ih => apply analz_spies_mono; simp_all

-- Unicity for NS1: nonce NA identifies agents A and B
theorem unique_NA { h : ns_public evs } :
  (Crypt (pubEK B) ⦃Nonce NA, Agent A⦄ ∈ parts (spies evs) →
  (Crypt (pubEK B') ⦃Nonce NA, Agent A'⦄ ∈ parts (spies evs) →
  (Nonce NA ∉ analz (spies evs) →
  A = A' ∧ B = B'))) := by
  induction h with
  | Nil => simp[spies, knows]
  | Fake _ a a_ih =>
    apply Fake_parts_sing at a; intro h₁ h₂ h₃;
    apply analz_spies_mono_neg at h₃;
    simp [*] at *
    apply Fake_parts_sing_helper (h := a) at h₁
    apply Fake_parts_sing_helper (h := a) at h₂
    simp_all
  | NS1 =>
    intro h₁ h₂ _; simp [*] at *
    expand_parts_element at h₁; expand_parts_element at h₂
    grind [ analz_insert_mono_neg, parts_knows_Spy_subset_used_neg ]
  | NS2 => intro _ _ h₃; apply analz_insert_mono_neg at h₃; simp_all
  | NS3 => intro _ _ h₃; apply analz_insert_mono_neg at h₃; simp_all;

-- Spy does not see the nonce sent in NS1 if A and B are secure
theorem Spy_not_see_NA { h : ns_public evs  }
  { not_bad_A : A ∉ bad }
  { not_bad_B : B ∉ bad } :
  Says A B (Crypt (pubEK B) ⦃Nonce NA, Agent A⦄) ∈ evs →
  Nonce NA ∉ analz (spies evs) := by
  intro h₁ h₄
  induction h with
  | Nil => simp_all
  | Fake _ a =>
    have _ := Spy_in_bad; apply Fake_analz_insert at a; apply a at h₄; simp_all
  | NS1 _ a a_ih =>
    simp_all; rcases h₁ with (_ | h)
    · simp_all; apply a; aapply analz_knows_Spy_subset_used
    · apply analz_insert_Crypt_subset at h₄; simp at h₄; cases h₄ <;> simp_all
      apply Says_imp_used at h; apply used_parts_subset_parts at h
      simp_all[Set.subset_def]
  | NS2 _ _ a a_ih =>
    simp [*] at *; have _ := h₄; have c := h₁
    apply Says_imp_parts_knows_Spy at h₁
    have d := h₁
    expand_parts_element at d
    apply analz_insert_Crypt_subset at h₄; simp at h₄; rcases h₄ with (h |h |h)
    <;> simp [*] at *;
    · apply a_ih at c; have _ := c; apply Says_imp_parts_knows_Spy at a
      apply unique_NA at h₁; apply h₁ at a; apply a at c; all_goals simp_all
    · grind[parts_knows_Spy_subset_used]
  | NS3 =>
    apply analz_insert_Crypt_subset at h₄; simp[*] at h₄;
    grind [Says_imp_parts_knows_Spy, no_nonce_NS1_NS2]

-- Authentication for `A`: if she receives message 2 and has used `NA` to start a run, then `B` has sent message 2.
theorem A_trusts_NS2 {h : ns_public evs }
  { not_bad_A : A ∉ bad }
  { not_bad_B : B ∉ bad } :
  Says A B (Crypt (pubEK B) ⦃Nonce NA, Agent A⦄) ∈ evs →
  Says B' A (Crypt (pubEK B) ⦃Nonce NA, Nonce NB, Agent B⦄) ∈ evs →
  Says B A (Crypt (pubEK B) ⦃Nonce NA, Nonce NB, Agent B⦄) ∈ 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 snsNA := h₁; apply Spy_not_see_NA at snsNA <;> try assumption
    apply analz_spies_mono_neg at snsNA
    simp [*] at *
    cases h₁
    · simp_all[Spy_in_bad]
    · apply Fake_parts_sing at a;
      apply Fake_parts_sing_helper (h := a) at h₂; simp at h₂
      grind [analz_subset_parts]
    · aapply ns_public.Fake
  | NS1 =>
    simp [*] at *; expand_parts_element at h₂
    grind[parts_knows_Spy_subset_used_neg]
  | NS2 =>
    simp [*] at *
    grind [ Spy_not_see_NA, Says_imp_parts_knows_Spy, unique_NA ]
  | NS3 => simp_all;

-- If the encrypted message appears then it originated with Alice in `NS1`
lemma B_trusts_NS1 { h : ns_public evs} :
  Crypt (pubEK B) ⦃Nonce NA, Agent A⦄ ∈ parts (spies evs) →
  Nonce NA ∉ analz (spies evs) →
  Says A B (Crypt (pubEK B) ⦃Nonce NA, Agent A⦄) ∈ evs
:= by
  intro h₁ h₂
  induction h with
  | Nil => simp[spies, knows] at h₁
  | Fake _ a a_ih =>
    apply analz_spies_mono_neg at h₂
    simp at h₁; apply Fake_parts_sing at a;
    apply Fake_parts_sing_helper (h := a) at h₁; simp_all
  | NS1 => apply analz_spies_mono_neg at h₂; simp_all; cases h₁ <;> simp_all
  | NS2 => apply analz_spies_mono_neg at h₂; simp_all;
  | NS3 => apply analz_spies_mono_neg at h₂; simp_all;

-- Authenticity Properties obtained from `NS2`

-- Unicity for `NS2`: nonce `NB` identifies nonce `NA` and agent `A`
theorem unique_NB { h : ns_public evs } :
  (Crypt (pubEK A) ⦃Nonce NA, Nonce NB, Agent B⦄ ∈ parts (spies evs) →
  (Crypt (pubEK A') ⦃Nonce NA', Nonce NB, Agent B'⦄ ∈ parts (spies evs) →
  (Nonce NB ∉ analz (spies evs) →
  A = A' ∧ NA = NA' ∧ B = B'))) := by
  -- Proof closely follows that of unique_NA
  induction h with
  | Nil => aesop (add norm spies, norm knows, safe analz_insertI)
  | Fake _ a a_ih =>
    apply Fake_parts_sing at a; intro h₁ h₂ h₃; simp [*] at *
    apply Fake_parts_sing_helper (h := a) at h₁;
    apply Fake_parts_sing_helper (h := a) at h₂; simp [*] at *
    apply analz_insert_mono_neg at h₃
    grind[analz_subset_parts]
  | NS1 => intro _ _ h₃; apply analz_spies_mono_neg at h₃; simp_all
  | NS2 =>
    intro h₁ h₂ _; simp [*] at *
    expand_parts_element at h₁
    expand_parts_element at h₂
    grind[analz_insert_mono_neg, parts_knows_Spy_subset_used]
  | NS3 => intro _ _ _; simp_all; grind[analz_insert_mono_neg]

-- `NB` remains secret
theorem Spy_not_see_NB { h : ns_public evs }
  { not_bad_A : A ∉ bad }
  { not_bad_B : B ∉ bad } :
  Says B A (Crypt (pubEK A) ⦃Nonce NA, Nonce NB, Agent B⦄) ∈ evs →
  Nonce NB ∉ 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[Spy_in_bad];
  | NS1 =>
    simp [*] at *
    apply analz_insert_Crypt_subset at h₄; simp at h₄
    have h₂ := h₁; apply Says_imp_parts_knows_Spy at h₂
    expand_parts_element at h₂
    grind[parts_knows_Spy_subset_used]
  | NS2 =>
    simp [*] at *
    have _ := h₄
    apply analz_insert_Crypt_subset at h₄; simp at h₄
    rcases h₁ with (_ | h₁)
    · simp_all; grind[parts_knows_Spy_subset_used, analz_subset_parts]
    · have _ := h₁; apply Says_imp_parts_knows_Spy at h₁
      expand_parts_element at h₁
      grind[
        parts_knows_Spy_subset_used,
        Says_imp_parts_knows_Spy,
        no_nonce_NS1_NS2
      ];
  | NS3 =>
    simp at h₁; simp[analz_insert_Crypt_element] at h₄;
    rcases h₄ with (⟨_, _⟩ | ⟨_, _⟩) <;> simp_all
    grind[Says_imp_parts_knows_Spy, unique_NB]

-- Authentication for `B`: if he receives message 3 and has used `NB` in message
-- 2, then `A` has sent message 3.
theorem B_trusts_NS3 { h : ns_public evs }
  { not_bad_A : A ∉ bad }
  { not_bad_B : B ∉ bad } :
  Says B A (Crypt (pubEK A) ⦃Nonce NA, Nonce NB, Agent B⦄) ∈ evs →
  Says A' B (Crypt (pubEK B) (Nonce NB)) ∈ evs →
  Says A B (Crypt (pubEK B) (Nonce NB)) ∈ evs
:= by
  intro h₁ h₂
  apply Says_imp_parts_knows_Spy at h₂
  induction h with
  | Nil => simp_all
  | Fake _ a a_ih =>
    simp [*] at *
    apply Fake_parts_sing at a
    apply Fake_parts_sing_helper (h := a) at h₂; simp at h₂
    grind [ Spy_in_bad, analz_subset_parts, Spy_not_see_NB ]
  | NS1 => simp_all
  | NS2 =>
    simp [*] at *; expand_parts_element at h₂
    grind[ parts_knows_Spy_subset_used ];
  | NS3 =>
    simp [*] at *;
    grind [Spy_not_see_NB, Says_imp_parts_knows_Spy, unique_NB]

-- Overall guarantee for `B`

-- If NS3 has been sent and the nonce NB agrees with the nonce B joined with NA,
-- then A initiated the run using NA
theorem B_trusts_protocol { h : ns_public evs  }
  { not_bad_A : A ∉ bad }
  { not_bad_B : B ∉ bad } :
  Crypt (pubEK B) (Nonce NB) ∈ parts (spies evs) →
  Says B A (Crypt (pubEK A) ⦃Nonce NA, Nonce NB, Agent B⦄) ∈ evs →
  Says A B (Crypt (pubEK B) ⦃Nonce NA, Agent A⦄) ∈ evs := by
  intro h₁ h₂
  induction h with
  | Nil => simp_all
  | Fake _ a a_ih =>
    simp [*] at *
    apply Fake_parts_sing at a
    apply Fake_parts_sing_helper (h := a) at h₁;
    expand_parts_element at h₁
    grind[Spy_in_bad, analz_subset_parts, Spy_not_see_NB]
  | NS1 => simp_all
  | NS2 _ nonce_not_used a a_ih =>
    simp [*] at *; expand_parts_element at h₁
    grind[parts_knows_Spy_subset_used];
  | NS3 =>
    simp [*] at *
    grind[Spy_not_see_NB, Says_imp_parts_knows_Spy, unique_NB ]

end NS_Public