#pragma once
#include "rosenpass/prf.mpv"
#include "rosenpass/handshake_state.mpv"
#include "rosenpass/protocol.mpv"

DECL_SETUP(key)

DECL_SETUP(kem_sk)
#define kem_pk_tmpl kem_sk_tmpl
letfun setup_kem_pk(sks:kem_pk_tmpl) =
  kem_pub(setup_kem_sk(sks)).
#define SETUP_KEM_PAIR(sk, pk, setup) \
  sk <- setup_kem_sk(setup);          \
  pk <- kem_pub(sk);

// TODO: Model use of multiple shared keys
// TODO: Hide shk inside the kem abstraction?
fun biscuit_key(kem_sk) : key [private].

#define SETUP_SERVER(biscuit_key, sk, pk, setup) \
  SETUP_KEM_PAIR(sk, pk, setup)         \
  biscuit_key <- biscuit_key(sk);

#define SETUP_HANDSHAKE_STATE()         \
  SETUP_SERVER(biscuit_key, sskm, spkm, Ssskm)   \
  psk  <- setup_key(Spsk);             \
  spkt <- setup_kem_pk(Sspkt);

type seed.
fun rng_key(seed) : key.
fun rng_kem_sk(seed) : kem_sk.
DECL_SETUP(seed)
#define RNG_KEM_PAIR(sk, pk, setup)        \
  sk <- rng_kem_sk(setup_seed(setup)); \
  pk <- kem_pub(sk);

event ConsumeSeed(Atom, seed, Atom).
const Sptr, Epti, Spti, Eski:Atom.

restriction s:seed, p1:Atom, p2:Atom, ad1:Atom, ad2:Atom;
  event(ConsumeSeed(p1, s, ad1)) && event(ConsumeSeed(p2, s, ad2))
    ==> p1 = p2 && ad1 = ad2.

#include "rosenpass/responder.macro"
fun Cinit_conf(kem_sk_tmpl, key_tmpl, kem_pk_tmpl, InitConf_t) : Atom [data].
CK_EV(  event OskOinit_conf(key, key). )
MTX_EV( event ICRjct(InitConf_t, key, kem_sk, kem_pk). )
SES_EV( event ResponderSession(InitConf_t, key). )
event ConsumeBiscuit(Atom, kem_sk, kem_pk, Atom).

let Oinit_conf_inner(Ssskm:kem_sk_tmpl, Spsk:key_tmpl, Sspkt:kem_sk_tmpl, ic:InitConf_t) =
#if RANDOMIZED_CALL_IDS
  new call:Atom;
#else
  call <- Cinit_conf(Ssskm, Spsk, Sspkt, ic);
#endif

  SETUP_HANDSHAKE_STATE()

  eski <- kem_sk0;
  epki <- kem_pk0;
  let try_ = (
    INITCONF_CONSUME()
    event ConsumeBiscuit(biscuit_no, sskm, spkt, call);
    CK_EV(  event OskOinit_conf(ck_rh, osk); )
    SES_EV( event ResponderSession(ic, osk); )
    0
  ) in (
    NOP
  ) else (
#if MESSAGE_TRANSMISSION_EVENTS
    MTX_EV( event ICRjct(ic, psk, sskr, spki) )
#else
    0
#endif
  ).
  
let Oinit_conf() = 
  in(C, Cinit_conf(Ssskm, Spsk, Sspkt, ic));
  Oinit_conf_inner(Ssskm, Spsk, Sspkt, ic).

restriction biscuit_no:Atom, sskm:kem_sk, spkr:kem_pk, ad1:Atom, ad2:Atom;
  event(ConsumeBiscuit(biscuit_no, sskm, spkr, ad1)) && event(ConsumeBiscuit(biscuit_no, sskm, spkr, ad2))
    ==> ad1 = ad2.

// TODO: Restriction biscuit no invalidation

#include "rosenpass/initiator.macro"
fun Cresp_hello(RespHello_t) : Atom [data].
CK_EV(  event OskOresp_hello(key, key, key). )
MTX_EV( event RHRjct(RespHello_t, key, kem_sk, kem_pk). )
MTX_EV( event ICSent(RespHello_t, InitConf_t, key, kem_sk, kem_pk). )
SES_EV( event InitiatorSession(RespHello_t, key). )
let Oresp_hello(HS_DECL_ARGS, C_in:channel) =
  in(C_in, Cresp_hello(RespHello(sidr, =sidi, ecti, scti, biscuit, auth)));
  rh <- RespHello(sidr, sidi, ecti, scti, biscuit, auth);
  /* try */ let ic = (
    ck_ini <- ck;
    RESPHELLO_CONSUME()
    ck_ih <- ck;
    INITCONF_PRODUCE()
    CK_EV (event OskOresp_hello(ck_ini, ck_ih, osk); ) // TODO: Queries testing that there is no duplication
    MTX_EV( event ICSent(rh, ic, psk, sski, spkr); )
    SES_EV( event InitiatorSession(rh, osk); )
    ic
  /* success */ ) in (
    out(C_in, Envelope(create_mac(spkt, IC2b(ic)), IC2b(ic)))
  /* fail */ ) else (
#if MESSAGE_TRANSMISSION_EVENTS
    event RHRjct(rh, psk, sski, spkr)
#else
    0
#endif
  ).

// TODO: Restriction: Biscuit no invalidation

#include "rosenpass/responder.macro"
fun Cinit_hello(SessionId, Atom, kem_sk_tmpl, key_tmpl, kem_pk_tmpl, seed_tmpl, seed_tmpl, InitHello_t) : Atom [data].
CK_EV(  event OskOinit_hello(key, key, key, kem_sk, kem_pk, kem_pk, key, key). )
MTX_EV( event IHRjct(InitHello_t, key, kem_sk, kem_pk). )
MTX_EV( event RHSent(InitHello_t, RespHello_t, key, kem_sk, kem_pk). )
event ConsumeSidr(SessionId, Atom).
event ConsumeBn(Atom, kem_sk, kem_pk, Atom).

let Oinit_hello_inner(sidm:SessionId, biscuit_no:Atom, Ssskm:kem_sk_tmpl, Spsk:key_tmpl, Sspkt: kem_sk_tmpl, Septi: seed_tmpl, Sspti: seed_tmpl, ih: InitHello_t, C_out:channel) = 
#if RANDOMIZED_CALL_IDS
  new call:Atom;
#else
  call <- Cinit_hello(sidr, biscuit_no, Ssskm, Spsk, Sspkt, Septi, Sspti, ih);
#endif
  // TODO: This is ugly
  let InitHello(sidi, epki, sctr, pidiC, auth) = ih in

  SETUP_HANDSHAKE_STATE()

  eski <- kem_sk0;

  event ConsumeBn(biscuit_no, sskm, spkt, call);
  event ConsumeSidr(sidr, call);

  epti <- rng_key(setup_seed(Septi)); // RHR4
  spti <- rng_key(setup_seed(Sspti)); // RHR5
  event ConsumeSeed(Epti, setup_seed(Septi), call);
  event ConsumeSeed(Spti, setup_seed(Sspti), call);

  let rh = (
    INITHELLO_CONSUME()
    ck_ini <- ck;
    RESPHELLO_PRODUCE()
    CK_EV( event OskOinit_hello(ck_ini, ck, psk, sskr, spki, epki, epti, spti); ) // TODO: Queries testing that there is no duplication
    MTX_EV( event RHSent(ih, rh, psk, sskr, spki); )
    rh
  /* success */ ) in (
    out(C_out, Envelope(create_mac(spkt, RH2b(rh)), RH2b(rh)))

  /* fail */ ) else (
#if MESSAGE_TRANSMISSION_EVENTS
    event IHRjct(ih, psk, sskr, spki)
#else
    0
#endif
  ).

let Oinit_hello() = 
  in(C, Cinit_hello(sidr, biscuit_no, Ssskm, Spsk, Sspkt, Septi, Sspti, ih));
  Oinit_hello_inner(sidr, biscuit_no, Ssskm, Spsk, Sspkt, Septi, Sspti, ih, C).

restriction sid:SessionId, ad1:Atom, ad2:Atom;
  event(ConsumeSidr(sid, ad1)) && event(ConsumeSidr(sid, ad2))
    ==> ad1 = ad2.

restriction biscuit_no:Atom, sskm:kem_sk, spkr:kem_pk, ad1:Atom, ad2:Atom;
  event(ConsumeBn(biscuit_no, sskm, spkr, ad1)) && event(ConsumeBn(biscuit_no, sskm, spkr, ad2))
    ==> ad1 = ad2.

// TODO: Restriction: Attacker may not reuse session ids
// TODO: Restriction: Attacker may not reuse biscuit no

#include "rosenpass/initiator.macro"
fun Cinitiator(SessionId, kem_sk_tmpl, key_tmpl, kem_pk_tmpl, seed_tmpl, seed_tmpl) : Atom [data].
CK_EV( event OskOinitiator_ck(key). )
CK_EV( event OskOinitiator(key, key, kem_sk, kem_pk, key). )
MTX_EV( event IHSent(InitHello_t, key, kem_sk, kem_pk). )
event ConsumeSidi(SessionId, Atom).

let Oinitiator_inner(sidi: SessionId, Ssskm: kem_sk_tmpl, Spsk: key_tmpl, Sspkt: kem_sk_tmpl, Seski: seed_tmpl, Ssptr: seed_tmpl, C_out:channel) =
  #if RANDOMIZED_CALL_IDS
    new call:Atom;
  #else
    call <- Cinitiator(sidi, Ssskm, Spsk, Sspkt, Seski, Ssptr);
  #endif

    SETUP_HANDSHAKE_STATE()

    sidr <- sid0;

    RNG_KEM_PAIR(eski, epki, Seski) // IHI3
    sptr <- rng_key(setup_seed(Ssptr)); // IHI5
    event ConsumeSidi(sidi, call);
    event ConsumeSeed(Sptr, setup_seed(Ssptr), call);
    event ConsumeSeed(Eski, setup_seed(Seski), call);

    INITHELLO_PRODUCE()
    CK_EV( event OskOinitiator_ck(ck); ) 
    CK_EV( event OskOinitiator(ck, psk, sski, spkr, sptr); )
    MTX_EV( event IHSent(ih, psk, sski, spkr); )
    out(C_out, Envelope(create_mac(spkt, IH2b(ih)), IH2b(ih)));
    Oresp_hello(HS_PASS_ARGS, C_out).

let Oinitiator() =
  in(C, Cinitiator(sidi, Ssskm, Spsk, Sspkt, Seski, Ssptr));
  Oinitiator_inner(sidi, Ssskm, Spsk, Sspkt, Seski, Ssptr, C).

restriction sid:SessionId, ad1:Atom, ad2:Atom;
  event(ConsumeSidi(sid, ad1)) && event(ConsumeSidi(sid, ad2))
    ==> ad1 = ad2.

// TODO: Should this be modeled without an oracle
fun Creveal_kem_pk(kem_sk_tmpl) : Atom [data].
event RevealPk(kem_pk).
let Oreveal_kem_pk =
  in(C, Creveal_kem_pk(Spk));
  pk <- setup_kem_pk(Spk);
  event RevealPk(pk);
  out(C, pk).

let rosenpass_main() = 0
  | !Oreveal_kem_pk
  | REP(INITIATOR_BOUND, Oinitiator)
  | REP(RESPONDER_BOUND, Oinit_hello)
  | REP(RESPONDER_BOUND, Oinit_conf).

nounif v:seed_prec;   attacker(prepare_seed(trusted_seed( v )))/6217[hypothesis].
nounif v:seed;        attacker(prepare_seed( v ))/6216[hypothesis].
nounif v:seed;        attacker(rng_kem_sk( v ))/6215[hypothesis].
nounif v:seed;        attacker(rng_key( v ))/6214[hypothesis].
nounif v:key_prec;    attacker(prepare_key(trusted_key( v )))/6213[hypothesis].
nounif v:kem_sk_prec; attacker(prepare_kem_sk(trusted_kem_sk( v )))/6212[hypothesis].
nounif v:key;         attacker(prepare_key( v ))/6211[hypothesis].
nounif v:kem_sk;      attacker(prepare_kem_sk( v ))/6210[hypothesis].
nounif Spk:kem_sk_tmpl;
  attacker(Creveal_kem_pk(Spk))/6110[conclusion].
nounif sid:SessionId, Ssskm:kem_sk_tmpl, Spsk:key_tmpl, Sspkt:kem_sk_tmpl, Seski:seed_tmpl, Ssptr:seed_tmpl;
  attacker(Cinitiator( *sid, *Ssskm, *Spsk, *Sspkt, *Seski, *Ssptr ))/6109[conclusion].
nounif sid:SessionId, biscuit_no:Atom, Ssskm:kem_sk_tmpl, Spsk:key_tmpl, Sspkt:kem_sk_tmpl, Septi:seed_tmpl, Sspti:seed_tmpl, ih:InitHello_t;
  attacker(Cinit_hello( *sid, *biscuit_no, *Ssskm, *Spsk, *Sspkt, *Septi, *Sspti, *ih ))/6108[conclusion].
nounif rh:RespHello_t;
  attacker(Cresp_hello( *rh ))/6107[conclusion].
nounif Ssskm:kem_sk_tmpl, Spsk:key_tmpl, Sspkt:kem_sk_tmpl, ic:InitConf_t;
  attacker(Cinit_conf( *Ssskm, *Spsk, *Sspkt, *ic ))/6106[conclusion].
