#include "config.mpv"

#define CHAINING_KEY_EVENTS 1
#define MESSAGE_TRANSMISSION_EVENTS 1
#define SESSION_START_EVENTS 0
#define RANDOMIZED_CALL_IDS 0
#undef FULL_MODEL
#undef SIMPLE_MODEL
#define SIMPLE_MODEL 1

#include "prelude/basic.mpv"
#include "crypto/key.mpv"
#include "rosenpass/oracles.mpv"
#include "crypto/kem.mpv"

#define INITIATOR_TEST
#define NEW_TRUSTED_SEED(name)            \
  new MCAT(name, _secret_seed):seed_prec; \
  name <- make_trusted_seed(MCAT(name, _secret_seed)); \

free D:channel [private].
free secure_biscuit_no:Atom [private].
free secure_sidi,secure_sidr:SessionId [private].
free secure_psk:key [private].
free initiator1, initiator2:kem_sk_prec.
free responder1, responder2:kem_sk_prec.

let secure_init_hello(initiator: kem_sk_tmpl, sidi : SessionId, psk: key_tmpl, responder: kem_sk_tmpl) =
  NEW_TRUSTED_SEED(seski_trusted_seed)
  NEW_TRUSTED_SEED(ssptr_trusted_seed)
  Oinitiator_inner(sidi, initiator, psk, responder, seski_trusted_seed, ssptr_trusted_seed, D).

let secure_resp_hello(initiator: kem_sk_tmpl, responder: kem_sk_tmpl, sidr:SessionId, sidi:SessionId, biscuit_no:Atom, psk:key_tmpl) =
  in(D, Envelope(k, IH2b(InitHello(=sidi, epki, sctr, pidiC, auth))));
  ih <- InitHello(sidi, epki, sctr, pidiC, auth);
  NEW_TRUSTED_SEED(septi_trusted_seed)
  NEW_TRUSTED_SEED(sspti_trusted_seed)
  Oinit_hello_inner(sidr, biscuit_no, responder, psk, initiator, septi_trusted_seed, sspti_trusted_seed, ih, D).

let secure_init_conf(initiator: kem_sk_tmpl, responder: kem_sk_tmpl, psk:key_tmpl, sidi:SessionId, sidr:SessionId) =
  in(D, Envelope(k3, IC2b(InitConf(=sidi, =sidr, biscuit, auth3))));
  ic <- InitConf(sidi,sidr,biscuit, auth3);
  NEW_TRUSTED_SEED(seski_trusted_seed)
  NEW_TRUSTED_SEED(ssptr_trusted_seed)
  Oinit_conf_inner(initiator, psk, responder, ic).

let secure_communication(initiator: kem_sk_tmpl, responder:kem_sk_tmpl) =
  secure_key <- prepare_key(secure_psk);
  (!secure_init_hello(initiator, secure_sidi, secure_key, responder))
  | !secure_resp_hello(initiator, responder, secure_sidr, secure_sidi, secure_biscuit_no, secure_key)
  | !(secure_init_conf(initiator, responder, secure_key, secure_sidi, secure_sidr)).

let pipeChannel(D:channel, C:channel) =
  in(D, b:bits);
  out(C, b).

fun kem_private(kem_pk): kem_sk
  reduc forall sk_tmpl:kem_sk;
    kem_private(kem_pub(sk_tmpl)) = sk_tmpl[private].

let secretCommunication() =
#ifdef INITIATOR_TEST
  initiator_pk <- choice[setup_kem_pk(make_trusted_kem_sk(initiator1)), setup_kem_pk(make_trusted_kem_sk(initiator2))];
  initiator_seed <- prepare_kem_sk(kem_private(initiator_pk)); 
#else
  initiator_seed <- prepare_kem_sk(trusted_kem_sk(initiator1));
#endif
#ifdef RESPONDER_TEST
  responder_pk <- choice[setup_kem_pk(make_trusted_kem_sk(responder1)), setup_kem_pk(make_trusted_kem_sk(responder2))];
  responder_seed <- prepare_kem_sk(kem_private(responder_pk));
#else
  responder_seed <- prepare_kem_sk(trusted_kem_sk(responder1));
#endif
  secure_communication(initiator_seed, responder_seed) | !pipeChannel(D, C).

let reveal_pks() =
  out(C, setup_kem_pk(make_trusted_kem_sk(responder1)));
  out(C, setup_kem_pk(make_trusted_kem_sk(responder2)));
  out(C, setup_kem_pk(make_trusted_kem_sk(initiator1)));
  out(C, setup_kem_pk(make_trusted_kem_sk(initiator2))).

let rosenpass_main2() = 
  REP(INITIATOR_BOUND, Oinitiator)
  | REP(RESPONDER_BOUND, Oinit_hello)
  | REP(RESPONDER_BOUND, Oinit_conf).

let identity_hiding_main() = 
  0 | reveal_pks() |  rosenpass_main2() | phase 1; secretCommunication().

let main = identity_hiding_main.