val initialize_handshake_state :
     #nc : config
  -> protocol_name : hashable nc
  -> prologue : hashable nc
  -> s : option (keypair nc)
  -> e : option (keypair nc)
  -> rs : option (public_key nc)
  -> re : option (public_key nc)
  -> psk : option preshared_key
  -> Tot (handshake_state nc)
val initialize_handshake_state :
     #nc : config
  -> protocol_name : hashable nc
  -> prologue : hashable nc
  -> s : option (keypair nc)
  -> e : option (keypair nc)
  -> rs : option (public_key nc)
  -> re : option (public_key nc)
  -> psk : option preshared_key
  -> Tot (handshake_state nc)
let initialize_handshake_state protocol_name prologue s e rs re psk =
  (* initialize the symmetric state with the protocol name *)
  let sym_st1 = initialize_symmetric protocol_name in
  (* mix hash the prologue *)
  let sym_st2 = mix_hash prologue sym_st1 in
  (* return state *)
  {
    sym_state = sym_st2;
    static = s;
    ephemeral = e;
    remote_static = rs;
    remote_ephemeral = re;
    preshared = psk;
  }
module Spec.Noise.Base.Definitions

open FStar.Mul

open Lib.IntTypes
open Lib.RawIntTypes
open Lib.ByteSequence
open Lib.Sequence

open Spec.Noise.CryptoPrimitives

#set-options "--z3rlimit 15 --fuel 0 --ifuel 1"

/// Small helper for the proofs


(*** Cipher state *)

let initialize_key k = { k = k; n = 0; }

let has_key st =
  match st.k with
  | None -> false
  | Some _ -> true

let set_nonce n st = ({ st with n = n; })

let encrypt_with_ad nc aad plain st =
  match st.k with
  | None -> Res (plain, st)
  | Some k' ->
    if st.n >= saturated_nonce_value then Fail Saturated_nonce else
    if length aad > aead_max_input nc then Fail Input_size else
    if length plain > aead_max_input nc then Fail Input_size else
    let cipher = aead_encrypt nc k' st.n aad plain in
    Res (cipher, {st with n = st.n + 1})

let decrypt_with_ad nc aad c st =
  match st.k with
  | None -> Res (c, st)
  | Some k' ->
    if st.n >= saturated_nonce_value then Fail Saturated_nonce else
    if length aad > aead_max_input nc then Fail Input_size else
    if length c < aead_tag_size then Fail Input_size else
    if length c > aead_max_input nc + aead_tag_size  then Fail Input_size else
    let r = aead_decrypt nc k' st.n aad c in
    match r with
    | None -> Fail Decryption
    | Some decrypted ->
      Res (decrypted, {st with n = st.n + 1})

(*** Symmetric state *)

let initialize_symmetric #nc pname =
  let pname_len : size_nat = length pname in
  let h : hash nc =
  begin
    if pname_len <= hash_size nc
    then
      begin
      let pad_size = (hash_size nc - pname_len) in
      let zeros = create pad_size (u8 0) in
      concat #uint8 #pname_len #pad_size pname zeros
      end
    else do_hash nc pname
  end in
  { c_state = initialize_key None; ck = h; h = h }


(** Facility function not in the Noise technical paper: if the hash size is 64,
  * truncates a hash to so that it can be used as a key *)
val key_from_hash : #nc:config -> hash nc -> aead_key
let key_from_hash #nc h = if hash_size nc = 64 then Seq.slice h 0 32 else h

let mix_key #nc key st =
  let ck, temp_k = kdf2 st.ck key in
  let temp_k' = key_from_hash temp_k in
  let c_st = initialize_key (Some temp_k') in
  { st with ck = ck; c_state = c_st }

let mix_hash #nc data st =
  let cs = concat #uint8 #(hash_size nc) #(length data) st.h data in
  { st with h = do_hash nc cs }

let mix_key_and_hash #nc key st =
  let ck, temp_h, temp_k = kdf3 st.ck key in
  let temp_k' = key_from_hash temp_k in
  let st1 = { st with ck = ck } in
  (**) hash_max_input_norm_lem nc;
  let st2 = mix_hash temp_h st1 in
  let c_st = initialize_key (Some temp_k') in
  { st2 with c_state = c_st }

let encrypt_and_hash #nc plain st =
  let temp = encrypt_with_ad nc st.h plain st.c_state in
  match temp with
  | Fail x -> Fail x
  | Res (cipher, c_st') -> 
    let st1 = { st with c_state = c_st' } in
    let st2 = mix_hash cipher st1 in
    Res (cipher, st2)

let decrypt_and_hash #nc cipher st =
  let temp = decrypt_with_ad nc st.h cipher st.c_state in
  match temp with
  | Fail x -> Fail x
  | Res (plain, c_st') ->
    let st1 = { st with c_state = c_st' } in
    let st2 = mix_hash cipher st1 in
    Res (plain, st2)

let split st =  
  let temp_k1, temp_k2 = kdf2_no_key st.ck in
  let temp_k1_1 = key_from_hash temp_k1 in
  let temp_k2_1 = key_from_hash temp_k2 in
  let c1 = initialize_key (Some temp_k1_1) in
  let c2 = initialize_key (Some temp_k2_1) in
  (c1, c2)

(*** Handshake state *)
/
Spec.Noise.CryptoPrimitives.fsti.checked prims.fst.checked Lib.Sequence.fsti.checked Lib.RawIntTypes.fsti.checked Lib.IntTypes.fsti.checked Lib.ByteSequence.fsti.checked FStar.Seq.fst.checked FStar.Pervasives.Native.fst.checked FStar.Pervasives.fsti.checked FStar.Mul.fst.checked FStar.List.Tot.fst.checked
Spec.Noise.Base.Definitions.fst

    protocol_name: Spec.Noise.CryptoPrimitives.hashable nc ->
    prologue: Spec.Noise.CryptoPrimitives.hashable nc ->
    s: FStar.Pervasives.Native.option (Spec.Noise.CryptoPrimitives.keypair nc) ->
    e: FStar.Pervasives.Native.option (Spec.Noise.CryptoPrimitives.keypair nc) ->
    rs: FStar.Pervasives.Native.option (Spec.Noise.CryptoPrimitives.public_key nc) ->
    re: FStar.Pervasives.Native.option (Spec.Noise.CryptoPrimitives.public_key nc) ->
    psk: FStar.Pervasives.Native.option Spec.Noise.CryptoPrimitives.preshared_key
  -> Spec.Noise.Base.Definitions.handshake_state nc
Spec.Noise.CryptoPrimitives.config Spec.Noise.CryptoPrimitives.hashable FStar.Pervasives.Native.option Spec.Noise.CryptoPrimitives.keypair Spec.Noise.CryptoPrimitives.public_key Spec.Noise.CryptoPrimitives.preshared_key Spec.Noise.Base.Definitions.Mkhandshake_state Spec.Noise.Base.Definitions.symmetric_state Spec.Noise.Base.Definitions.mix_hash Spec.Noise.Base.Definitions.initialize_symmetric Spec.Noise.Base.Definitions.handshake_state
let initialize_handshake_state protocol_name prologue s e rs re psk =
  
let sym_st1 = initialize_symmetric protocol_name in
let sym_st2 = mix_hash prologue sym_st1 in
{
  sym_state = sym_st2;
  static = s;
  ephemeral = e;
  remote_static = rs;
  remote_ephemeral = re;
  preshared = psk
}