(*
  Sender Keys protocol; proving secrecy, authentication, PCS, and PFS
  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, 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 P2P MASTER + SENDER KEYS 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 3: SENDER KEYS SESSION B -> A =================== *)

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

  let ca2 = hkdf(concat3(concat0(khash(ra1)), dh(gtb2, ta1))) in
  let ak2 = khash(ca2) in
  let ak2_auth = hkdf(concat0(ak2)) in 

  if checkmac(ak2_auth, concat2(x2, gtb2), x2_mac) = okay then (

    let mak2_enc = hkdf(concat0(ak2)) in 

    let (MPK_B: pkey, mrb0: key) = sdec(mak2_enc, x2) in
    event recvOE2(MPK_B, mrb0, gao, gta1);

    (* =================== PHASE 4: SENDER KEYS 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_sig = sign(MSK_A, mx1) in
    event sendME1(m1, mra1, MPK_A);
    out(c, (mx1, mx1_sig));

    (* =================== PHASE 5: SENDER KEYS MSG B -> A =================== *)

    in(c, (mx2: bitstring, mx2_sig: bitstring));
    let mrb1 = khash(mrb0) in
    let mbk1_auth = hkdf(concat0(mrb1)) in
    let mbk1_enc = hkdf(concat0(mrb1)) in 
    if checksign(MPK_B, mx2, mx2_sig) = okay then 
      let m2 = sdec(mbk1_enc, mx2) in 
      event recvME2(m2, mrb1, MPK_B);

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

      event compromiseOSKA(OSK_A);
      out(c, OSK_A)
  ).

let PeerB(OSK_B: skey, OPK_B: pkey, OPK_A: pkey, 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 P2P MASTER + SENDER KEYS 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 3: SENDER KEYS SESH B -> A =================== *)

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

    new tb2: skey;
    let gtb2 = pk(tb2) in 
    let rb2 = hkdf(concat3(concat0(khash(rb1)), dh(gta1, tb2))) in 
    let cb2 = hkdf(concat3(concat0(khash(rb1)), dh(gta1, tb2))) in 

    let bk2 = khash(cb2) in
    let bk2_auth = hkdf(concat0(bk2)) in 
    let bk2_enc = hkdf(concat0(bk2)) in 

    new mrb0: key; (* mrb0 := sender keys ratchet *)

    let x2 = senc(bk2_enc, (MPK_B, mrb0)) in 
    let x2_mac = mac(bk2_auth, concat2(x2, gtb2)) in
    event sendOE2(MPK_B, mrb0, gao, gta1);

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


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

    in(c, (mx1: 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, mx1, mx1_sig) = okay then (
      let m1 = sdec(mak1_enc, mx1) in 
      event recvME1(m1, mra1, MPK_A);

      phase 3;

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

      (* =================== PHASE 5: SENDER KEYS MSG B -> A =================== *)

      let mrb1 = khash(mrb0) in
      let mbk1_auth = hkdf(concat0(mrb1)) in
      let mbk1_enc = hkdf(concat0(mrb1)) in 
      
      let mx2 = senc(mbk1_enc, m2) in
      let mx2_sig = sign(MSK_B, mx2) in
      event sendME2(m2, mrb1, MPK_B);
      out(c, (mx2, mx2_sig))

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

  )).

process
  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);
  (*
  out(a, MPK_A); 
  out(a, MPK_B); 
  *)

  new fib1: skey;
  new fib2: skey;

  new fib3: skey;
  new fib4: skey;
  let (kO_A: skey, kM_A: skey) = choice [(OSK_A, MSK_A), (fib1, fib2) ] in
  let (kO_B: skey, kM_B: skey) = choice [(OSK_B, MSK_B), (fib1, fib2) ] in

  ( (!PeerA(kO_A, pk(kO_A), OPK_B, kM_A, pk(kM_A))) |
    (!PeerB(OSK_B, OPK_B, OPK_A, MSK_B, MPK_B)))