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

-- 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, grind =]
theorem Spy_see_priEK {h : ns_public evs} :
  (Key (priEK A) ∈ parts (spies evs)) ↔ A ∈ bad := by
  constructor
  · induction h with
    | Nil => simp[priEK]
    | 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 grind

-- It is impossible to re-use a nonce in both NS1 and NS2, provided the nonce is
-- secret
@[grind! .]
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 at h₂
  | Fake _ h =>
    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_all; grind
  | NS1 =>
    simp [*] at *
    expand_parts_element at h₁; expand_parts_element at h₂;
    grind
  | NS2 =>
    simp [*] at *
    expand_parts_element at h₂;
    grind
  | NS3 => apply analz_spies_mono; simp_all

-- Unicity for NS1: nonce NA identifies agents A and B
@[grind! .]
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
  intro h₁ h₂ h₃
  induction h with
  | Nil => simp_all
  | Fake _ a a_ih =>
    apply Fake_parts_sing at a;
    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 =>
    simp [*] at *; expand_parts_element at h₁; expand_parts_element at h₂; grind
  | NS2 => simp_all; grind
  | NS3 => simp_all; grind

-- Spy does not see the nonce sent in NS1 if A and B are secure
@[grind! .]
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 => apply Fake_analz_insert at a; apply a at h₄; simp_all
  | NS1 _ a =>
    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
  | NS3 => apply analz_insert_Crypt_subset at h₄; simp[*] at h₄; grind

-- 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
    · apply Fake_parts_sing at a;
      apply Fake_parts_sing_helper (h := a) at h₂; simp at h₂
      grind
    · aapply ns_public.Fake
  | NS1 => simp [*] at *; expand_parts_element at h₂; grind
  | NS2 => simp [*] at *; grind
  | 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 at h₁
  | Fake _ a =>
    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; grind
  | 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`
@[grind! .]
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
  intro h₁ h₂ h₃
  induction h with
  | Nil => aesop (add safe analz_insertI)
  | Fake _ a =>
    apply Fake_parts_sing at a; 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
  | NS1 => apply analz_spies_mono_neg at h₃; simp_all
  | NS2 =>
    simp [*] at *; expand_parts_element at h₁; expand_parts_element at h₂; grind
  | NS3 => simp_all; grind

-- `NB` remains secret
@[grind! .]
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;
  | NS1 =>
    simp [*] at *
    apply analz_insert_Crypt_subset at h₄; simp at h₄
    have _ := h₁; apply Says_imp_parts_knows_Spy at h₁
    expand_parts_element at h₁
    grind
  | NS2 =>
    simp [*] at *
    have _ := h₄
    apply analz_insert_Crypt_subset at h₄; simp at h₄;
    rcases h₁ with (_ | h₁)
    · simp_all; grind
    · have _ := h₁; apply Says_imp_parts_knows_Spy at h₁;
      expand_parts_element at h₁; grind
  | NS3 => simp[analz_insert_Crypt_element] at h₄; simp [*] at *; grind

-- 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 =>
    simp [*] at *
    apply Fake_parts_sing at a
    apply Fake_parts_sing_helper (h := a) at h₂; simp at h₂
    grind
  | NS1 => simp_all
  | NS2 => simp [*] at *; expand_parts_element at h₂; grind
  | NS3 => simp [*] at *; grind

-- 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 =>
    simp [*] at *
    apply Fake_parts_sing at a
    apply Fake_parts_sing_helper (h := a) at h₁; expand_parts_element at h₁
    grind
  | NS1 => simp_all
  | NS2 => simp [*] at *; expand_parts_element at h₁; grind
  | NS3 => simp [*] at *; grind

end NS_Public