(*
  Olm + Megolm protocol
  Author: [redacted]
*)

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

(*
set selFun = Nounifset.   
set redundancyElim = best.
set simpEqAll = false. 
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 key encryption *)

type skey.
type pkey.

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

reduc forall m: bitstring, sk: skey; adec(aenc(m,pk(sk)),sk) = m.

(* 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): bitstring.
equation forall a: skey, b: skey; dh(pk(a), b) = dh(pk(b), a). (* symmetry of DH *)

(* the concat functions *)
fun khash(key): key.
fun hkdf(bitstring): key [data]. 

fun concat0(key): bitstring [data, typeConverter].
fun concat1(bitstring, pkey, pkey): bitstring [data].
reduc forall m1: bitstring, k1: pkey, k2: pkey; split1(concat1(m1, k1, k2)) = (m1, k1, k2). 

fun concat2(bitstring, pkey): bitstring [data].
reduc forall m1: bitstring, k1: pkey; split2(concat2(m1, k1)) = (m1, k1). 
                              
fun concat3(bitstring, bitstring): bitstring [data].
reduc forall m1: bitstring, m2: bitstring; split3(concat3(m1, m2)) = (m1, m2).
fun concat4(bitstring, bitstring, bitstring): bitstring [data].

reduc forall m1: bitstring, m2: bitstring, m3:bitstring; split4(concat4(m1, m2, m3)) = (m1, m2, m3).
                                              
equation forall k: key, a: skey, b: skey; hkdf(concat3(concat0(khash(k)), dh(pk(a), b))) = hkdf(concat3(concat0(khash(k)), dh(pk(b), a))).

(* the concats *)

(* events *)
event sendOE1(pkey, key, pkey, pkey).
event recvOE1(pkey, key, pkey, pkey).

event sendOE2(pkey, key, pkey, pkey).
event recvOE2(pkey, key, pkey, pkey).

event sendME1(bitstring, key, pkey).
event recvME1(bitstring, key, pkey).

event sendME2(bitstring, key, pkey).
event recvME2(bitstring, key, pkey).

event compromiseOSKA(skey).

event compromiseOSKB(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(OSK_A: skey, OPK_A: pkey, OPK_B: pkey, OLM_KEYS_STR:bitstring, OLM_ROOT_STR: bitstring, MSK_A: skey, MPK_A: pkey) =
  phase 1;

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

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

  (* in(c, gbo: pkey); *)
  in(c, (gbo: pkey, gbo_sig: bitstring));
  (* if checksign(OPK_B, rb(gbo), gbo_sig) = okay then ( *)

  (* =================== PHASE 2: DERIVE OLM MASTER + MEGOLM SESSION A -> B =================== *)
  
  let amaster = hkdf(concat4(dh(OPK_B, OSK_A), dh(gbo, OSK_A), dh(OPK_B, ao))) in
  let ra1 = hkdf(concat0(amaster)) in
  let ca1 = hkdf(concat0(amaster)) in

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

  let ak1 = khash(ca1) in
  let ak1_auth = hkdf(concat0(ak1)) in 
  let ak1_enc = hkdf(concat0(ak1)) in 

  new mra0: key; (* mra0 := megolm ratchet *)

  let x1 = senc(ak1_enc, (MPK_A, mra0)) in
  let x1_mac = mac(ak1_auth, concat1(x1, gao, gta1)) in
  event sendOE1(MPK_A, mra0, gao, gta1);

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

  (* =================== PHASE 4: MEGOLM MSG A -> B =================== *)

  let mra1 = khash(mra0) in 
  let mak1_auth = hkdf(concat0(mra1)) in
  let mak1_enc = hkdf(concat0(mra1)) in

  let mx1 = senc(mak1_enc, m1) in
  let mx1_mac = mac(mak1_auth, mx1) in
  let mx1_sig = sign(MSK_A, concat3(mx1, mx1_mac)) in
  event sendME1(m1, mra1, MPK_A);
  out(c, (mx1, mx1_mac, mx1_sig));

  phase 3; (* =================== PHASE 6: PCS =================== *)

  (* out(c, PK_A); *)

  0.

let PeerB(OSK_B: skey, OPK_B: pkey, OPK_A: pkey, OLM_KEYS_STR:bitstring, OLM_ROOT_STR: bitstring, MSK_B: skey, MPK_B: pkey) =
  new bo: skey;

  let gbo = pk(bo) in
  let gbo_sig = sign(OSK_B, rb(gbo)) in

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

  (* =================== PHASE 2: DERIVE OLM MASTER + MEGOLM SESSION A -> B =================== *)

  (* 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 = hkdf(concat4(dh(OPK_A, OSK_B), dh(OPK_A, bo), dh(gao, OSK_B))) in
  let rb1 = hkdf(concat0(bmaster)) in  
  let cb1 = hkdf(concat0(bmaster)) in  

  let bk1 = khash(cb1) in
  let bk1_auth = hkdf(concat0(bk1)) in
  let bk1_enc = hkdf(concat0(bk1)) in
  if checkmac(bk1_auth, concat1(x1, gao, gta1), x1_mac) = okay then 
  (
    let (MPK_A: pkey, mra0: key) = sdec(bk1_enc, x1) in 
    event recvOE1(MPK_A, mra0, gao, gta1);

   (* =================== PHASE 4: MEGOLM MSG A -> B =================== *)

    in(c, (mx1: bitstring, mx1_mac: bitstring, mx1_sig: bitstring));

    let mra1 = khash(mra0) in
    let mak1_auth = hkdf(concat0(mra1)) in 
    let mak1_enc = hkdf(concat0(mra1)) in 

    if checksign(MPK_A, concat3(mx1, mx1_mac), mx1_sig) = okay then (
      if checkmac(mak1_auth, mx1, mx1_mac) = okay then (
        let m1 = sdec(mak1_enc, mx1) in 
        event recvME1(m1, mra1, MPK_A);

        (* =================== PHASE 6: PCS =================== *)

        phase 3;

        event compromiseOSKB(OSK_B);
        out(c, OSK_B)

  ))).



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

(* 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(compromiseOSKB(sk)).  

(* post-compromise security. even if the secret key is compromised, message two remains secret 
*)

query sk1: skey, sk2: skey; (event(compromiseOSKB(sk1)) && attacker(m1)) ==> false.

(* auth *)
query k1: pkey, rk: key, k2: pkey, k3: pkey; inj-event(recvOE1(k1, rk, k2, k3)) ==> inj-event(sendOE1(k1, rk, k2, k3)).
query m: bitstring, rk: key, k1: pkey; inj-event(recvME1(m, rk, k1)) ==> inj-event(sendME1(m, rk, k1)).

(* secrecy *)
query attacker(m1). (* expected to be false *)

(* reachability: all expected to be false *) 
query k1: pkey, rk: key, k2: pkey, k3: pkey; event(sendOE1(k1,rk,k2,k3)).
query k1: pkey, rk: key, k2: pkey, k3: pkey; event(recvOE1(k1,rk,k2,k3)).

query m: bitstring, rk: key, k1: pkey; event(sendME1(m, rk, k1)).
query m: bitstring, rk: key, k1: pkey; event(recvME1(m, rk, k1)).

process
  new OLM_KEYS_STR:bitstring; out(a, OLM_KEYS_STR); 
  new OLM_ROOT_STR:bitstring; out(a, OLM_ROOT_STR); 
  new OSK_A: skey; let OPK_A = pk(OSK_A) in 
  new OSK_B: skey; let OPK_B = pk(OSK_B) in

  new MSK_A: skey; let MPK_A = pk(MSK_A) in  
  new MSK_B: skey; let MPK_B = pk(MSK_B) in 
  out(a, OPK_A);
  out(a, OPK_B);
  (*
   do not publish sender keys material
  out(a, MPK_A); 
  out(a, MPK_B); 
  *)
  event start();
  ( (!PeerA(OSK_A, OPK_A, OPK_B, OLM_KEYS_STR, OLM_ROOT_STR, MSK_A, MPK_A)) |
    (!PeerB(OSK_B, OPK_B, OPK_A, OLM_KEYS_STR, OLM_ROOT_STR, MSK_B, MPK_B)))