(*
  Olm 3DH+Double Ratchet without olm signing; 
  Author: [redacted]
*)

free m1: bitstring [private].
free m2: bitstring [private].

set simpEqAll = false. 
set selFun = Nounifset.   
set redundancyElim = best.
set redundantHypElim = true.  
set simplifyProcess = true.
set stopTerm = false.

free c: channel.
free a: channel. (* channel for the attacker *)
free p: channel [private]. (* For the distribution of public keys with integrity and authenticity - verification happens out of band. This is a standard assumption. *)

(* Symmetric key encryption *)

type key.

fun senc(key, bitstring): bitstring.
reduc forall m: bitstring, k: key; sdec(k, senc(k,m)) = m.

(* Asymmetric keys *)

type skey.
type pkey.

fun rb(pkey): bitstring.
fun pk(skey): pkey.

(* Digital signatures *)

fun sign(skey, bitstring): bitstring.
fun okay():bitstring.
reduc forall m: bitstring, sk: skey; checksign(pk(sk), m, sign(sk, m)) = okay.

(* MACs *)
fun mac(key, bitstring): bitstring.
reduc forall k: key, m: bitstring; checkmac(k, m, mac(k, m)) = okay.

(* Diffie-Hellman *)
(* DH -> Public^Private *)

fun dh(pkey, skey): key.
equation forall a: skey, b: skey; dh(pk(a), b) = dh(pk(b), a). (* symmetry of DH *)

(* the concat functions *)
fun hkdf1(key, key, key): key [data].

fun khash(key): key.

fun hkdf2_dev1(key): key.
fun hkdf2_dev2(key): key.
letfun hkdf2(k: key) =
  (hkdf2_dev1(k), hkdf2_dev2(k)).

fun hkdf3_dev1(key, bitstring): key.
fun hkdf3_dev2(key, bitstring): key.
letfun hkdf3(k: key, b: bitstring) =  
  (hkdf3_dev1(k, b), hkdf3_dev2(k, b)).

fun hkdf4_dev1(key, key): key.
fun hkdf4_dev2(key, key): key.
letfun hkdf4(k1: key, k2: key) =  
  (hkdf4_dev1(k1, k2), hkdf4_dev2(k1, k2)).

(* the concats *)

fun concat1(bitstring, pkey, pkey): bitstring [data].
fun concat2(bitstring, pkey): bitstring [data].

(* events *)
event sendE1(bitstring, key, pkey, pkey).
event recvE1(bitstring, key, pkey, pkey).

event sendE2(bitstring, key, pkey, pkey).
event recvE2(bitstring, key, pkey, pkey).

event compromiseSKA(skey).

event compromiseSKB(skey).

event start().

free tag_oe1: bitstring [private].
free tag_oe2: bitstring [private].
free tag_me1: bitstring [private].
free tag_me2: bitstring [private].
free tag_b_eph: bitstring [private].

let PeerA(SK_A: skey, PK_A: pkey, PK_B: pkey, OLM_KEYS_STR:bitstring, OLM_ROOT_STR: bitstring) =
  phase 1;

  new ao: skey;
  let gao = pk(ao) in

  (* generate amaster and enc msg (PHASE 1) *)

  in(c, (gbo: pkey));
  
  let amaster = hkdf1(dh(PK_B, SK_A), dh(gbo, SK_A), dh(PK_B, ao)) in
  let (ra1: key, ca1: key) = hkdf3(amaster, OLM_ROOT_STR) in (* derive the root and chain key *)

  new ta1: skey;
  let gta1 = pk(ta1) in 

  let mak1 = khash(ca1) in
  let (mak1_auth: key, mak1_enc: key) = hkdf2(mak1) in

  let x1 = senc(mak1_enc, m1) in 
  let x1_mac = mac(mak1_auth, concat1(x1, gao, gta1)) in
  event sendE1(m1, mak1, gao, gta1);

  out(c, (x1, x1_mac, gao, gta1));

  (* second stage: now, decrypt the received message from bob *)
  in(c, (x2: bitstring, x2_mac: bitstring, gtb2: pkey));

  let (ra2: key, ca2: key) = hkdf4(ra1, dh(gtb2, ta1)) in
  let mak2 = khash(ca2) in
  let (mak2_auth: key, mak2_enc: key) = hkdf2(mak2) in

  if checkmac(mak2_auth, concat2(x2, gtb2), x2_mac) = okay then
  ( 
    let m2 = sdec(mak2_enc, x2) in
    event recvE2(m2, mak2, gta1, gtb2);
    phase 2
  ).

let PeerB(SK_B: skey, PK_B: pkey, PK_A: pkey, OLM_KEYS_STR:bitstring, OLM_ROOT_STR: bitstring) =
  new bo: skey;

  let gbo = pk(bo) in

  out(c, (gbo)); 
  phase 1;

   (* first stage: derive bmaster, verfiy a's msgs, decrypt prekey message, reply *)

  in(c, (x1: bitstring, x1_mac: bitstring, gao: pkey, gta1: pkey)); 

  let bmaster = hkdf1(dh(PK_A, SK_B), dh(PK_A, bo), dh(gao, SK_B)) in
  let (rb1: key, cb1: key) = hkdf3(bmaster, OLM_ROOT_STR) in (* derive the root and chain key *)

  let mbk1 = khash(cb1) in
  let (mbk1_auth: key, mbk1_enc: key) = hkdf2(mbk1) in
  if checkmac(mbk1_auth, concat1(x1, gao, gta1), x1_mac) = okay then 
  (
    let m1 = sdec(mbk1_enc, x1) in 
    event recvE1(m1, mbk1, gao, gta1);

    new tb2: skey;
    let gtb2 = pk(tb2) in 

    let (rb2: key, cb2: key) = hkdf4(rb1, dh(gta1, tb2)) in
    let mbk2 = khash(cb2) in
    let (mbk2_auth: key, mbk2_enc: key) = hkdf2(mbk2) in

    let x2 = senc(mbk2_enc, m2) in 
    let x2_mac = mac(mbk2_auth, concat2(x2, gtb2)) in
    event sendE2(m2, mbk2, gta1, gtb2);

    out(c, (x2, x2_mac, gtb2));

    phase 2;
    event compromiseSKB(SK_B);
    out(c, SK_B)
  ).

query event(start()). (* reachable from all possible executions *)

(* pre-compromise security, aka forward secrecy. the only way
 m1 can be compromised is if alice's sk is compromised 

 NOTE: if signed, this is trivially true since m1 is never compromised
*)
query sk: skey; attacker(m1) ==> event(compromiseSKB(sk)).  

(* post-compromise security. even if the secret key is compromised, message two remains secret 
*)
query sk: skey; (event(compromiseSKB(sk)) && attacker(m2)) ==> false.

(* auth *)
query m: bitstring, rk: key, k1: pkey, k2: pkey; inj-event(recvE1(m, rk, k1, k2)) ==> inj-event(sendE1(m, rk, k1, k2)).
query m: bitstring, rk: key, k1: pkey, k2: pkey, k3: pkey, k4: pkey; inj-event(recvE2(m, rk, k1, k2)) ==> inj-event(sendE2(m, rk, k1, k2)).

(* secrecy *)

query attacker(m1). 
query attacker(m2). 

(* reachability *) 
query m: bitstring, rk: key, k1: pkey, k2: pkey; event(recvE1(m, rk, k1, k2)). (* reachable from all executions *)
query m: bitstring, rk: key, k1: pkey, k2: pkey; event(recvE2(m, rk, k1, k2)). (* reachable from all executions *)
query m: bitstring, rk: key, k1: pkey, k2: pkey; event(sendE1(m, rk, k1, k2)). (* reachable from all executions *)
query m: bitstring, rk: key, k1: pkey, k2: pkey; event(sendE2(m, rk, k1, k2)). (* rechable from all executions *)

process
  new OLM_KEYS_STR:bitstring; out(a, OLM_KEYS_STR); 
  new OLM_ROOT_STR:bitstring; out(a, OLM_ROOT_STR); 
  new SK_A: skey; let PK_A = pk(SK_A) in 
  new SK_B: skey; let PK_B = pk(SK_B) in
  out(a, PK_A);
  out(a, PK_B);
  event start();
  ( (!PeerA(SK_A, PK_A, PK_B, OLM_KEYS_STR, OLM_ROOT_STR)) |
    (!PeerB(SK_B, PK_B, PK_A, OLM_KEYS_STR, OLM_ROOT_STR)))