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

-- Lemmata for some very specific recurring cases in the following proof
omit [InvKey] [Bad] in
lemma Fake_parts_sing_helper {A B : Set Msg}
{ h : A ⊆ B } :
X ∈ A ∨ h₁ → X ∈ B ∨ h₁
:= by
  intro h; cases h <;> try simp_all
  left; aapply h
  
-- 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]; intro h
      rcases h with (((⟨B, bad, h⟩ | ⟨B, bad, h⟩) | ⟨B, h⟩) | ⟨B, h⟩) <;>
      try (apply injective_publicKey at h; simp_all)
      all_goals (apply publicKey_neq_privateKey at h; contradiction)
    | 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; right
      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 ((_ | _) | _) <;>
      try simp_all
      all_goals (aapply ih <;> aapply analz_subset_parts)
  | NS1 _ nonce_not_used =>
    apply parts_knows_Spy_subset_used_neg at nonce_not_used;
    simp[spies] at h₁; expand_parts_element at h₁;
    simp[spies] at h₂; expand_parts_element at h₂;
    apply analz_mono; apply Set.subset_insert
    cases h₂ <;> simp_all
  | NS2 _ nonce_not_used =>
    apply parts_knows_Spy_subset_used_neg at nonce_not_used;
    simp[spies] at h₁; expand_parts_element at h₁;
    simp[spies] at h₂; expand_parts_element at h₂;
    apply analz_mono; apply Set.subset_insert
    cases h₁ <;> simp_all
  | NS3 _ _ _ a_ih => simp at h₁; simp at h₂; apply analz_mono
                      apply Set.subset_insert; aapply a_ih

@[simp]
lemma injective_pubEK_helper:
  ( pubEK A = pubEK B ∧ h) ↔ ( A = B ∧ h )
:= by
  constructor
  · intro h₁
    rcases h₁ with ⟨e, _⟩
    apply injective_publicKey at e
    aapply And.intro; simp_all
  · intro h₁; 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 => 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 h₁; apply Fake_parts_sing_helper (h := a) at h₁; simp at h₁
    simp at h₂; apply Fake_parts_sing_helper (h := a) at h₂; simp at h₂
    simp at h₃;
    rcases h₁ with ((_ | _) | _) <;> 
    rcases h₂ with ((_ | _) | _) <;>
    try (
      apply False.elim; apply h₃; apply analz_mono; aapply Set.subset_insert
      tauto
    )
    all_goals (aapply a_ih <;> try aapply analz_subset_parts
               all_goals (
                 intro _; apply h₃; aapply analz_mono; aapply Set.subset_insert
               ))
  | NS1 _ nonce_not_used a_ih =>
    intro h₁ h₂ h₃
    simp at h₁; expand_parts_element at h₁
    simp at h₂; expand_parts_element at h₂
    apply parts_knows_Spy_subset_used_neg at nonce_not_used
    cases h₁ <;> cases h₂ <;> simp_all
    -- TODO make an analz_mono_neg lemma for these cases
    aapply a_ih; intro _; apply h₃; apply_rules[analz_mono, Set.subset_insert]
  | NS2 _ _ _ a_ih => intro h₁ h₂ h₃; simp_all; apply a_ih; intro h; apply h₃
                      apply_rules [analz_mono, Set.subset_insert]
  | NS3 _ _ _ a_ih => intro h₁ h₂ h₃; simp at h₁; simp at h₂; aapply a_ih
                      intro h; apply h₃
                      apply_rules [analz_mono, Set.subset_insert]

-- 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 a_ih => 
    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 => 
    cases h₁ with | tail _ b =>
      have _ := h₄
      simp at h₄; apply analz_insert_Crypt_subset at h₄
      simp at h₄; rcases h₄ with ( h | h | h)
      · simp at a_ih; have c := b; apply a_ih at c; rw[h] at b;
        have _ := c; rw[h] at c;
        apply Says_imp_parts_knows_Spy at b
        apply Says_imp_parts_knows_Spy at a
        apply unique_NA at b; apply b at a; apply a at c; simp_all
        assumption
      · rw [h] at b
        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₂ a_ih => 
    cases h₁ with | tail _ b =>
      have _ := h₄
      simp at h₄; apply analz_insert_Crypt_subset at h₄
      simp at h₄; rcases h₄ with ( h | h )
      · have _ := b; have _ := a₁; have _ := a₂
        rw[h] at b; apply Says_imp_parts_knows_Spy at b
        apply Says_imp_parts_knows_Spy at a₂
        aapply a_ih; apply no_nonce_NS1_NS2
        · assumption
        · rw[h]; exact a₂
        · rw[h]; exact b
      · aapply a_ih
      
-- 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
    simp at h₁; simp at h₂;
    cases h₁
    · have _ := Spy_in_bad; simp_all
    · right; apply Fake_parts_sing at a;
      apply Fake_parts_sing_helper (h := a) at h₂; simp at h₂
      rcases h₂ with ((_ | _) | _) <;> aapply a_ih
      · aapply analz_subset_parts
      · apply False.elim; apply snsNA; apply analz_spies_mono; tauto;
    · aapply ns_public.Fake
  | NS1 _ a a_ih =>
    simp at h₂; expand_parts_element at h₂
    apply parts_knows_Spy_subset_used_neg at a; cases h₁ <;> simp_all
    aapply a_ih 
  | NS2 _ _ a a_ih =>
    simp at h₁; have b := h₁; have snsNA := h₁
    apply Spy_not_see_NA at snsNA <;> try assumption
    simp at h₂; 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 at h₁; simp at h₂; right; aapply a_ih
  
-- 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 =>
    simp at h₁; apply Fake_parts_sing at a;
    apply Fake_parts_sing_helper (h := a) at h₁; simp at h₁
    rcases h₁ with ((h₁ | h₁ )| h₁); 
    · right; aapply a_ih; aapply analz_subset_parts; aapply analz_spies_mono_neg
    · apply False.elim; apply h₂; apply analz_spies_mono; simp_all
    · right; aapply a_ih; aapply analz_spies_mono_neg
  | NS1 _ _ a_ih => simp at h₁; cases h₁
                    · simp_all
                    · right; aapply a_ih; aapply analz_spies_mono_neg
  | NS2 _ _ _ a_ih => simp at h₁; right; aapply a_ih; aapply analz_spies_mono_neg
  | NS3 _ _ _ a_ih => simp at h₁; right; aapply a_ih; aapply analz_spies_mono_neg
  
-- 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 h₁; apply Fake_parts_sing_helper (h := a) at h₁; simp at h₁;
    simp at h₂; apply Fake_parts_sing_helper (h := a) at h₂; simp at h₂;
    apply analz_spies_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 h₁; expand_parts_element at h₁
    simp at h₂; expand_parts_element at h₂
    apply analz_spies_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;

-- `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 h₁
    simp[spies, knows] at h₄; apply analz_insert_Crypt_subset at h₄; simp at h₄
    apply parts_knows_Spy_subset_used_neg at nonce_not_used
    cases h₄ with
    | inl e => apply Says_imp_parts_knows_Spy at h₁;
               expand_parts_element at h₁; simp_all
    | inr => aapply a_ih
  | NS2 _ not_used_NB a a_ih =>
    simp at h₁;
    simp[spies, knows] at h₄;
    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 =>
    right; simp at h₁
    apply Fake_parts_sing at a
    simp at h₂; apply Fake_parts_sing_helper (h := a) at h₂; simp at h₂
    expand_parts_element at h₂;
    have _ := Spy_in_bad
    rcases h₁ with (h₁ | h₁) <;> rcases h₂ with ((h₂ | h₂) | h₂) <;> simp_all
    · aapply a_ih; aapply analz_subset_parts
    · apply Spy_not_see_NB at h₁ <;> simp_all
    · aapply a_ih
  | NS1 _ a a_ih => right; simp at h₂; simp at h₁; aapply a_ih;
  | NS2 _ nonce_not_used a a_ih =>
    right
    apply parts_knows_Spy_subset_used_neg at nonce_not_used;
    simp at h₂; expand_parts_element at h₂;
    simp at h₁; cases h₁ <;> simp_all; aapply a_ih
  | NS3 _ _ a₂ a_ih =>
    simp at h₁
    simp at h₂; 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

-- 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 =>
    right
    apply Fake_parts_sing at a
    simp at h₁; apply Fake_parts_sing_helper (h := a) at h₁;
    expand_parts_element at h₁
    have _ := Spy_in_bad
    simp at h₂; rcases h₂ with (_ | h₂) <;> simp_all
    rcases h₁ with (((_ |_ ) | _) | _)  <;> try (aapply a_ih)
    · aapply analz_subset_parts
    · apply Spy_not_see_NB at h₂ <;> simp_all
    · simp_all
  | NS1 _ a a_ih => simp_all
  | NS2 _ nonce_not_used a a_ih =>
    simp at h₁; simp at h₂;
    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 h₂
    simp at h₁; 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