import InductiveVerification.Public
set_option diagnostics true

-- 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
set_option trace.aesop true
@[simp]
theorem Spy_see_priEK {h : ns_public evs} :
  (Key (priEK A) ∈ parts (spies evs)) ↔ A ∈ bad := by
  constructor
  -- · induction h <;> aesop (add norm spies, norm knows, norm initState, norm pubEK, norm priEK, norm pubSK, norm priSK, norm injective_publicKey)
  · induction h with
    | Nil => simp[spies, knows, initState]; intro h; cases h with
             | inl h => cases h with
               | inl => aesop (add norm pubEK, norm pubSK, safe forward injective_publicKey)
               | inr h => simp[pubEK, priEK] at h; cases h with
                | intro _ h => apply publicKey_neq_privateKey at h; contradiction
             | inr h => simp[pubSK, priEK] at h; cases h with
               | intro _ h => apply publicKey_neq_privateKey at h; contradiction
    | Fake _ h ih => apply Fake_parts_sing at h
                     simp[spies, knows]
                     intro h₁; apply ih;
                     cases h₁ with
                      | inl h₁ => apply h at h₁; cases h₁ with 
                        | inl h₁ => cases h₁; aapply analz_subset_parts
                        | inr => assumption
                      | inr => assumption
    | NS1 _ _ ih => aesop (add norm spies, norm knows)
    | NS2 _ _ _ ih => aesop (add norm spies, norm knows)
    | NS3 _ _ _ ih => rw[spies, knows]; simp; intro _; apply ih; grind
  · 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

-- Lammata for some very specific recurring cases in the following proof
lemma no_nonce_NS1_NS2_helper1
{h : synth (analz (spies evsf)) ⦃Nonce NA, Msg.Agent A⦄}
: Nonce NA ∈ analz (knows Agent.Spy evsf) := by
  cases h with
  | inj => aapply analz.fst
  | mpair n => cases n; assumption

lemma no_nonce_NS1_NS2_helper2
{ih : Crypt (pubEK C) ⦃NA', ⦃Nonce NA, Msg.Agent D⦄⦄ ∈ parts (spies evsf)
  → Crypt (pubEK B) ⦃Nonce NA, Msg.Agent A⦄ ∈ parts (spies evsf)
  → Nonce NA ∈ analz (spies evsf)}
{h : parts {X} ⊆ synth (analz (spies evsf)) ∪ parts (spies evsf)}
{h₁ : Crypt (pubEK C) ⦃NA', ⦃Nonce NA, Msg.Agent D⦄⦄ ∈ parts (knows Agent.Spy evsf)}
{h₂ : Crypt (pubEK B) ⦃Nonce NA, Msg.Agent A⦄ ∈ parts {X} ∨
      Crypt (pubEK B) ⦃Nonce NA, Msg.Agent A⦄ ∈ parts (knows Agent.Spy evsf)}
: Nonce NA ∈ analz (knows Agent.Spy evsf) := by
  apply ih at h₁; cases h₂ with
  | inl h₂ => apply h at h₂; cases h₂ with
    | inl h₂ => cases h₂ with
      | inj => apply h₁; aapply analz_subset_parts
      | crypt h₂ => aapply no_nonce_NS1_NS2_helper1
    | inr => aapply h₁
  | inr => aapply h₁

-- It is impossible to re-use a nonce in both NS1 and NS2, provided the nonce is secret
theorem no_nonce_NS1_NS2 { 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 => rw[spies, knows] at h₂; simp[initState] at h₂
  | Fake _ h ih => 
      apply Fake_parts_sing at h
      simp[spies, knows] at h₁
      apply analz_insert; right;
      cases h₁ with 
      | inl h₁ => simp_all; apply h at h₁; cases h₁ with
        | inl h₁ => cases h₁ with
          | inj h₁ => apply analz_subset_parts at h₁
                      aapply no_nonce_NS1_NS2_helper2
          | crypt h₁ => cases h₁ with
            | inj => apply analz.fst; aapply analz.snd
            | mpair _ h₁ => aapply no_nonce_NS1_NS2_helper1
        | inr h₁ => aapply no_nonce_NS1_NS2_helper2
      | inr h₁ => simp[spies, knows] at h₂; aapply no_nonce_NS1_NS2_helper2
  | NS1 =>
    simp[spies] at h₁; simp[spies] at h₂; cases h₂ with
    | inl h => rcases h with ⟨_ , ⟨n , _⟩⟩;
               apply parts.body at h₁; apply parts.snd at h₁; apply parts.fst at h₁
               apply parts_knows_Spy_subset_used at h₁; rw[n] at h₁; contradiction
    | inr => rw[spies, knows]; apply analz_mono; apply Set.subset_insert; simp_all
  | NS2 =>
    simp[spies] at h₁; simp[spies] at h₂; cases h₁ with
    | inl h => rcases h with ⟨_, ⟨_, ⟨n, _⟩⟩⟩
               apply parts.body at h₂; apply parts.fst at h₂
               apply parts_knows_Spy_subset_used at h₂; rw[n] at h₂; contradiction
    | inr => rw[spies, knows]; apply analz_mono; apply Set.subset_insert; simp_all
  | NS3 => rw[spies, knows]; apply analz_mono; apply Set.subset_insert; simp_all

lemma unique_NA_apply_ih {P : Prop}
{a_ih : Crypt (pubEK B) ⦃Nonce NA, Msg.Agent A⦄ ∈ parts (spies evsf)
  → Crypt (pubEK B') ⦃Nonce NA, Msg.Agent A'⦄ ∈ parts (spies evsf)
  → Nonce NA ∉ analz (spies evsf) → P}
{h₃ : Nonce NA ∉ analz (spies (Says Agent.Spy C X :: evsf))}
{h₁ : Crypt (pubEK B) ⦃Nonce NA, Msg.Agent A⦄ ∈ parts (spies evsf)}
{h₂ : Crypt (pubEK B') ⦃Nonce NA, Msg.Agent A'⦄ ∈ parts (spies evsf)}
: P := by
  simp[spies, knows] at h₃; apply Set.notMem_subset at h₃
  · aapply a_ih;
  · apply analz_mono; apply Set.subset_insert    

lemma unique_NA_contradict
{h₃ : Nonce NA ∉ analz (spies (Says Agent.Spy B X :: evsf))}
{h₂ : synth (analz (spies evsf)) ⦃Nonce NA, Msg.Agent A'⦄}
{P : Prop}
: P := by
  apply MPair_synth_analz.mp at h₂; rcases h₂ with ⟨⟨n⟩, m⟩;
  simp[spies, knows] at h₃; apply Set.notMem_subset at h₃
  · contradiction;
  · apply analz_mono; apply Set.subset_insert 

-- 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[spies, knows] at h₁; cases h₁ with
    | inl h₁ => apply a at h₁; cases h₁ with
      | inl h₁ => cases h₁ with
        | inj h₁ => simp[spies, knows] at h₂; cases h₂ with
          | inl h₂ => apply a at h₂; cases h₂ with
            | inl h₂ => cases h₂ with
              | inj h₂ => apply analz_subset_parts at h₁
                          apply analz_subset_parts at h₂
                          aapply unique_NA_apply_ih   
              | crypt h₂ => aapply unique_NA_contradict
            | inr h₂ => apply analz_subset_parts at h₁
                        aapply unique_NA_apply_ih
          | inr h₂ => apply analz_subset_parts at h₁
                      aapply unique_NA_apply_ih
        | crypt h₁ => aapply unique_NA_contradict
      | inr h₁ => simp[spies] at h₁; simp[spies, knows] at h₂; cases h₂ with
        | inl h₂ => apply a at h₂; cases h₂ with
          | inl h₂ => cases h₂ with
            | inj h₂ => apply analz_subset_parts at h₂
                        aapply unique_NA_apply_ih
            | crypt => aapply unique_NA_contradict
          | inr => aapply unique_NA_apply_ih
        | inr => aapply unique_NA_apply_ih
    | inr => simp[spies, knows] at h₂; cases h₂ with
      | inl h₂ => apply a at h₂; cases h₂ with
        | inl h₂ => cases h₂ with
          | inj h₂ => apply analz_subset_parts at h₂
                      aapply unique_NA_apply_ih
          | crypt => aapply unique_NA_contradict
        | inr => aapply unique_NA_apply_ih
      | inr => aapply unique_NA_apply_ih
  | NS1 _ _ a_ih => intro h₁ h₂ h₃; simp at h₁; cases h₁ with
    | inl => sorry
    | inr => simp at h₂; cases h₂ with
      | inl h₂ => simp at h₃; rw[analz_Crypt] at h₃
                  rcases h₂ with ⟨_, ⟨nonce_eq, _⟩⟩; rw[nonce_eq] at h₃; simp at h₃;
      | inr => aapply unique_NA_apply_ih;
  | NS2 => sorry
  | NS3 => sorry

-- Spy does not see the nonce sent in NS1 if A and B are secure
theorem Spy_not_see_NA { h : ns_public evs  }:
  Says A B (Crypt (pubEK B) ⦃Nonce NA, Agent A⦄) ∈ evs →
  A ∉ bad →
  B ∉ bad →
  Nonce NA ∉ analz (spies evs) := by
  intro h
  induction h <;> simp_all [analz_insertI, no_nonce_NS1_NS2]

-- 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  }:
  A ∉ bad →
  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
  induction h <;> simp_all [analz_insertI, no_nonce_NS1_NS2]

end NS_Public