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
      · left
    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]
    | Fake _ h ih =>
      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 h₁; apply parts_increasing; aapply 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 =>
      simp; apply analz_insert;
      apply Fake_parts_sing at h
      simp at h₁; apply Fake_parts_sing_helper (h := h) at h₁; simp at h₁
      simp at h₂; apply Fake_parts_sing_helper (h := h) at h₂; simp at h₂
      rcases h₁ with ((_ | _) | _) <;>
      rcases h₂ with ((_ | _) | _) <;>
      simp_all
      all_goals (right; aapply ih <;> aapply analz_subset_parts)
  | NS1 _ nonce_not_used =>
    apply analz_spies_mono
    simp [*] at *
    apply parts_knows_Spy_subset_used_neg at nonce_not_used;
    expand_parts_element at h₁; expand_parts_element at h₂;
    cases h₂ <;> simp_all
  | NS2 _ nonce_not_used =>
    apply analz_spies_mono
    simp [*] at *
    apply parts_knows_Spy_subset_used_neg at nonce_not_used;
    expand_parts_element at h₂;
    cases h₁ <;> simp_all[-Key.injEq]
  | 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 _ nonce_not_used a_ih =>
    intro h₁ h₂ h₃
    apply analz_insert_mono_neg at h₃
    simp [*] at *
    expand_parts_element at h₁
    expand_parts_element at h₂
    apply parts_knows_Spy_subset_used_neg at nonce_not_used
    cases h₁ <;> cases h₂ <;> simp_all
  | 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; cases h₁ with
    | inl   => simp_all; apply a; aapply analz_knows_Spy_subset_used
    | inr h => 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; apply a; apply h; simp
               · aapply a_ih
  | NS2 _ not_used_NB a a_ih =>
    simp at h₁
    have _ := h₄
    simp at h₄; apply analz_insert_Crypt_subset at h₄
    simp at h₄; rcases h₄ with ( h | h | h)
    · simp [*] at *; have c := h₁; apply a_ih at c;
      have _ := c;
      apply Says_imp_parts_knows_Spy at h₁
      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
    · simp_all
      apply not_used_NB; apply parts_knows_Spy_subset_used; apply parts.fst;
      apply parts.body; apply Says_imp_parts_knows_Spy; assumption
    · aapply a_ih
  | NS3 _ _ a₂ a_ih =>
    simp [*] at *
    have _ := h₄
    apply analz_insert_Crypt_subset at h₄; simp[*] at h₄;
    have _ := h₁; simp[*] at h₁; apply Says_imp_parts_knows_Spy at h₁
    apply Says_imp_parts_knows_Spy at a₂
    aapply a_ih; apply no_nonce_NS1_NS2 <;> try simp [*] <;> assumption

-- 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 a_ih =>
    have snsNA := h₁; apply Spy_not_see_NA at snsNA <;> try assumption
    apply analz_spies_mono_neg at snsNA
    simp [*] at *
    cases h₁
    · have _ := Spy_in_bad; simp_all
    · apply Fake_parts_sing at a;
      apply Fake_parts_sing_helper (h := a) at h₂; simp at h₂
      rcases h₂ with ((_ | _) | _) <;> (right; aapply a_ih)
      · aapply analz_subset_parts
      · tauto
    · aapply ns_public.Fake
  | NS1 _ a a_ih =>
    apply parts_knows_Spy_subset_used_neg at a;
    simp [*] at *; expand_parts_element at h₂; cases h₁ <;> simp_all
  | NS2 _ _ a a_ih =>
    simp [*] at *
    have snsNA := h₁; apply Spy_not_see_NA at snsNA <;> try assumption
    cases h₂ <;> simp_all
    apply Says_imp_parts_knows_Spy at a; apply unique_NA at a;
    apply Says_imp_parts_knows_Spy at h₁; apply a at h₁; all_goals simp_all
  | NS3 _ _ a a_ih => 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 _ _ a_ih =>
    apply analz_spies_mono_neg at h₂; simp_all; cases h₁ <;> simp_all
  | NS2 _ _ _ a_ih => apply analz_spies_mono_neg at h₂; simp_all;
  | NS3 _ _ _ a_ih => 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₃
    rcases h₁ with ((_ | _) | _) <;>
    rcases h₂ with ((_ | _) | _) <;>
    simp_all
    all_goals (aapply a_ih; repeat aapply analz_subset_parts)
  | NS1 _ _ a_ih => intro h₁ h₂ h₃; simp at h₁; simp at h₂; aapply a_ih
                    aapply analz_spies_mono_neg
  | NS2 _ nonce_not_used _ a_ih =>
    intro h₁ h₂ h₃; simp [*] at *
    expand_parts_element at h₁
    expand_parts_element at h₂
    apply analz_insert_mono_neg at h₃;
    apply parts_knows_Spy_subset_used_neg at nonce_not_used
    rcases h₁ with (_ | h₁) <;>
    rcases h₂ with (_ | h₂) <;> simp_all
  | NS3 _ _ _ a_ih =>
    intro h₁ h₂ h₃; apply analz_spies_mono_neg at h₃; simp_all[-Key.injEq]

-- `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 a_ih =>
    have _ := Spy_in_bad; apply Fake_analz_insert at a; apply a at h₄; simp_all;
  | NS1 _ nonce_not_used a_ih =>
    simp [*] at *
    apply analz_insert_Crypt_subset at h₄; simp at h₄
    apply parts_knows_Spy_subset_used_neg at nonce_not_used
    have h₂ := h₁; apply Says_imp_parts_knows_Spy at h₂
    expand_parts_element at h₂; simp_all
  | NS2 _ not_used_NB a a_ih =>
    simp [*] at *
    apply parts_knows_Spy_subset_used_neg at not_used_NB
    rcases h₁ with (_ | h₁)
    · simp_all; apply not_used_NB; aapply analz_subset_parts
    · apply analz_insert_Crypt_subset at h₄; simp at h₄; rcases h₄ with (_ |_ |_ )
      · aapply a_ih; apply Says_imp_parts_knows_Spy at a;
        apply Says_imp_parts_knows_Spy at h₁; simp_all; aapply no_nonce_NS1_NS2
      · apply Says_imp_parts_knows_Spy at h₁;
        expand_parts_element at h₁; simp_all
      · aapply a_ih
  | NS3 _ _ a a_ih =>
    simp at h₁; simp[analz_insert_Crypt_element] at h₄;
    rcases h₄ with (⟨_, _⟩ | ⟨_, _⟩) <;> simp_all
    apply Says_imp_parts_knows_Spy at a
    apply Says_imp_parts_knows_Spy at h₁; apply unique_NB at a
    apply a at h₁; apply h₁ at a_ih; simp_all; assumption

-- 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₂
    expand_parts_element at h₂;
    rcases h₁ with (_ | h₁) <;>
    rcases h₂ with ((h₂ | _) | _) <;> simp_all[Spy_in_bad]
    · apply analz_subset_parts at h₂; simp_all
    · apply Spy_not_see_NB at h₁ <;> simp_all
  | NS1 => simp_all
  | NS2 _ nonce_not_used =>
    simp [*] at *
    apply parts_knows_Spy_subset_used_neg at nonce_not_used;
    expand_parts_element at h₂; cases h₁ <;> simp_all
  | NS3 _ _ a₂ =>
    simp [*] at *;
    expand_parts_element at h₂; cases h₂ <;> simp_all
    have h₁c := h₁
    apply Spy_not_see_NB at h₁c
    apply Says_imp_parts_knows_Spy at h₁; apply unique_NB at h₁;
    apply Says_imp_parts_knows_Spy at a₂; apply h₁ at a₂
    all_goals simp_all

-- 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₁
    rcases h₂ with (_ | h₂) <;> simp_all[Spy_in_bad]
    rcases h₁ with (((_ |_ ) | _) | _)  <;> try simp_all
    · right; aapply a_ih; aapply analz_subset_parts
    · apply Spy_not_see_NB at h₂ <;> simp_all
  | NS1 => simp_all
  | NS2 _ nonce_not_used a a_ih =>
    simp [*] at *
    apply parts_knows_Spy_subset_used_neg at nonce_not_used;
    expand_parts_element at h₁; cases h₂ <;> simp_all
  | NS3 _ _ a₂ a_ih =>
    simp [*] at *
    expand_parts_element at h₁
    cases h₁ <;> simp_all
    have h₂c := h₂
    apply Spy_not_see_NB at h₂c
    apply Says_imp_parts_knows_Spy at h₂
    apply Says_imp_parts_knows_Spy at a₂
    apply unique_NB at h₂; apply h₂ at a₂
    apply a₂ at h₂c; all_goals simp_all

end NS_Public