open Global
open Base64
open Data
open Pi
open Crypto
open Prins

let debug_impl = debug "impl"

type principal = string

type mackey = bytes hkey
let get_mackey (a:principal) (b:principal) = Prins.get_mackey (cS a) (cS b)
let session = sha1 (utf8 (cS"(c_start,c_qcwrCommit),({q,c,w,r}Commit{q,c,w,r}),(w_start,w_qcwrCommit))(w_qcwrCommit,w_Fwd_qcwrCommit),({}Fwd{q,c,w,r}),(r_start,r_Fwd_qcwrCommit))"))


type vars = {
  c : principal ;
  w : principal ;
  r : principal ;
  q : int }
type hashes = {
  hc : bytes ;
  hw : bytes ;
  hr : bytes ;
  hq : bytes }
type macs = {
  wCommitqcwr : bytes ;
  rCommitqcwr : bytes ;
  rFwd : bytes }
type keys = {
  key_cw : bytes ;
  key_cr : bytes ;
  key_wc : bytes ;
  key_wr : bytes ;
  key_rc : bytes ;
  key_rw : bytes }
type header = {
  ts : int ;
  sid : bytes }
type store = 
  { vars : vars ; hashes : hashes ; macs : macs ; keys : keys ; header :  header}

let empty_store_c role starting = 
  {vars = { c = role ; w = "" ; r = "" ; q = 0 };
   hashes = { hc = utf8 (cS "") ; hw = utf8 (cS "") ; hr = utf8 (cS "") ; hq = utf8 (cS "") };
   macs = { wCommitqcwr = utf8 (cS "") ; rCommitqcwr = utf8 (cS "") ; rFwd = utf8 (cS "") };
   keys = { key_cw = utf8 (cS "") ;
  key_cr = utf8 (cS "") ;
  key_wc = utf8 (cS "") ;
  key_wr = utf8 (cS "") ;
  key_rc = utf8 (cS "") ;
  key_rw = utf8 (cS "") };
   header = { sid = if starting then sha1 (concat (mkNonce()) session) else utf8 (cS ""); ts = 0 } }

let empty_store_w role starting = 
  {vars = { c = "" ; w = role ; r = "" ; q = 0 };
   hashes = { hc = utf8 (cS "") ; hw = utf8 (cS "") ; hr = utf8 (cS "") ; hq = utf8 (cS "") };
   macs = { wCommitqcwr = utf8 (cS "") ; rCommitqcwr = utf8 (cS "") ; rFwd = utf8 (cS "") };
   keys = { key_cw = utf8 (cS "") ;
  key_cr = utf8 (cS "") ;
  key_wc = utf8 (cS "") ;
  key_wr = utf8 (cS "") ;
  key_rc = utf8 (cS "") ;
  key_rw = utf8 (cS "") };
   header = { sid = if starting then sha1 (concat (mkNonce()) session) else utf8 (cS ""); ts = 0 } }

let empty_store_r role starting = 
  {vars = { c = "" ; w = "" ; r = role ; q = 0 };
   hashes = { hc = utf8 (cS "") ; hw = utf8 (cS "") ; hr = utf8 (cS "") ; hq = utf8 (cS "") };
   macs = { wCommitqcwr = utf8 (cS "") ; rCommitqcwr = utf8 (cS "") ; rFwd = utf8 (cS "") };
   keys = { key_cw = utf8 (cS "") ;
  key_cr = utf8 (cS "") ;
  key_wc = utf8 (cS "") ;
  key_wr = utf8 (cS "") ;
  key_rc = utf8 (cS "") ;
  key_rw = utf8 (cS "") };
   header = { sid = if starting then sha1 (concat (mkNonce()) session) else utf8 (cS ""); ts = 0 } }

type preds = 
| Send_Commit of (principal * bytes * int* int * principal * principal * principal)
| Send_Fwd of (principal * bytes * int)
| Recv_Commit of (principal * bytes * int * int * principal * principal * principal)
| Recv_Fwd of (principal * bytes * int * int * principal * principal * principal)



(*******************************************)
(* Wired types, send and receive functions *)
(*******************************************)

let content_w_Fwd_qcwrCommit_ = fun ts store ->
  let empty_str = cS "" in
  let nil = utf8 empty_str in
  let hashes = nil in
  let nextstate = cS "w_Fwd_qcwrCommit" in
  let nextstate_string = utf8 nextstate in
  let payload = concat nextstate_string hashes in
  let ts_string = string_of_int ts in
  let ts_str = cS ts_string in
  let ts_mar = utf8 ts_str in
  let header = concat store.header.sid ts_mar in
  let content = concat header payload in
  content

let content_c_qcwrCommit_qcwr = fun ts store ->
  let empty_str = cS "" in
  let nil = utf8 empty_str in
  let hashes = nil in
  let hashes = concat store.hashes.hr hashes in
  let hashes = concat store.hashes.hw hashes in
  let hashes = concat store.hashes.hc hashes in
  let hashes = concat store.hashes.hq hashes in
  let nextstate = cS "c_qcwrCommit" in
  let nextstate_string = utf8 nextstate in
  let payload = concat nextstate_string hashes in
  let ts_string = string_of_int ts in
  let ts_str = cS ts_string in
  let ts_mar = utf8 ts_str in
  let header = concat store.header.sid ts_mar in
  let content = concat header payload in
  content

let mac_verify_w_c_qcwrCommit_qcwr = fun ts store k m ->
  let content = content_c_qcwrCommit_qcwr ts store in
  let cc = mac_verify k m content in
  let pc = pickle content in
  let _ = test_eq cc pc "MAC incorrect" in
  let up = unpickle cc in
  up

let mac_verify_r_w_Fwd_qcwrCommit_ = fun ts store k m ->
  let content = content_w_Fwd_qcwrCommit_ ts store in
  let cc = mac_verify k m content in
  let pc = pickle content in
  let _ = test_eq cc pc "MAC incorrect" in
  let up = unpickle cc in
  up

let mac_verify_r_c_qcwrCommit_qcwr = fun ts store k m ->
  let content = content_c_qcwrCommit_qcwr ts store in
  let cc = mac_verify k m content in
  let pc = pickle content in
  let _ = test_eq cc pc "MAC incorrect" in
  let up = unpickle cc in
  up



type wired_w_start =
  | Wired_in_w_start_of_c_qcwrCommit_qcwr__ of (int * principal * principal * principal) * store

type wired_r_start =
  | Wired_in_r_start_of_w_Fwd_qcwrCommit___c_qcwrCommit_qcwr__ of (int * principal * principal * principal) * store


let receiveWired_w_start : store -> wired_w_start = fun store ->
  let empty_str = cS "" in
  let nil = utf8 empty_str in
  let msg = precv  store.vars.w in
  let header,content = iconcat (ibase64 msg) in
  let sid,ts_mar = iconcat header in
  let ts_str = iutf8 ts_mar in
  let ts_string = iS ts_str in
  let ts = int_of_string ts_string in
  let () = check_cache store.vars.w "w" sid in
  let oldts = store.header.ts in
  let store = { 
      vars = store.vars;
      hashes = store.hashes;
      macs = store.macs; 
      keys = store.keys; 
      header = {sid = sid; ts = ts; }} in
  let tag,payload = iconcat content in
  match iS (iutf8 tag) with
  | "c_qcwrCommit_qcwr__" ->
  let variables,security = iconcat payload in
  let keys,protocol = iconcat security in
  let hashes,macs = iconcat protocol in
  (* Unmarshalling variables *)
  let encr_q,variables = iconcat variables in
  let mar_c,variables = iconcat variables in
  let string_c = iutf8 mar_c in
  let c = iS string_c in
  let hc = sha1 mar_c in
  let mar_w,variables = iconcat variables in
  let string_w = iutf8 mar_w in
  let w = iS string_w in
  let hw = sha1 mar_w in
  let mar_r,variables = iconcat variables in
  let string_r = iutf8 mar_r in
  let r = iS string_r in
  let hr = sha1 mar_r in
  let store = { 
      vars = {store.vars with c = c; w = w; r = r; };
      hashes = store.hashes;
      macs = store.macs; 
      keys = store.keys; 
      header = store.header} in
  (* Registering new keys *)
  let key_cw,keys = iconcat keys in
  let () = reg_keys (cS store.vars.c) (cS store.vars.w) key_cw in
  let key_cr,_ = iconcat keys in
  let () = reg_keys (cS store.vars.c) (cS store.vars.r) key_cr in
  (* Decrypting variables *)
  let keycw = get_symkey (cS store.vars.c) (cS store.vars.w) in
  let mar_q = unpickle (sym_decrypt keycw encr_q) in
  let string_q = iutf8 mar_q in
  let marred_q = iS string_q in
  let q = int_of_string marred_q in
  let hq = sha1 mar_q in
  let store = { 
      vars = {store.vars with q = q; };
      hashes = {store.hashes with hq = hq; hc = hc; hw = hw; hr = hr; };
      macs = store.macs; 
      keys = {store.keys with key_cw = key_cw; key_cr = key_cr; }; 
      header = store.header} in
  (* Unmarshalling hashes *)
  (* Unmarshalling MACs *)
  let wCommitqcwr,macs = iconcat macs in
  let rCommitqcwr,_ = iconcat macs in
  let store = { 
      vars = store.vars;
      hashes = store.hashes;
      macs = {store.macs with wCommitqcwr = wCommitqcwr; rCommitqcwr = rCommitqcwr; }; 
      keys = store.keys; 
      header = store.header} in
  (* Verification of a MAC from state c_qcwrCommit with variables qcwr*)
  let macheader,maccontent = iconcat store.macs.wCommitqcwr in
  let macsid,ts_mar = iconcat macheader in
  let ts_str = iutf8 ts_mar in
  let ts_string = iS ts_str in
  let ts = int_of_string ts_string in
  let _ = test_inf oldts ts "MAC verification" in
  let oldts = ts in
  let _ = test_eq macsid sid "MAC verification" in
  let mackeycw = get_mackey store.vars.c store.vars.w in
  let content = mac_verify_w_c_qcwrCommit_qcwr ts store mackeycw maccontent in
  let _ = test_eq oldts store.header.ts "Time-stamp verification" in
  (* Verification Ended *)
    let wired = Wired_in_w_start_of_c_qcwrCommit_qcwr__((q,c,w,r),store) in
    wired

let receiveWired_r_start : store -> wired_r_start = fun store ->
  let empty_str = cS "" in
  let nil = utf8 empty_str in
  let msg = precv  store.vars.r in
  let header,content = iconcat (ibase64 msg) in
  let sid,ts_mar = iconcat header in
  let ts_str = iutf8 ts_mar in
  let ts_string = iS ts_str in
  let ts = int_of_string ts_string in
  let () = check_cache store.vars.r "r" sid in
  let oldts = store.header.ts in
  let store = { 
      vars = store.vars;
      hashes = store.hashes;
      macs = store.macs; 
      keys = store.keys; 
      header = {sid = sid; ts = ts; }} in
  let tag,payload = iconcat content in
  match iS (iutf8 tag) with
  | "w_Fwd_qcwrCommit___c_qcwrCommit_qcwr__" ->
  let variables,security = iconcat payload in
  let keys,protocol = iconcat security in
  let hashes,macs = iconcat protocol in
  (* Unmarshalling variables *)
  let encr_q,variables = iconcat variables in
  let mar_c,variables = iconcat variables in
  let string_c = iutf8 mar_c in
  let c = iS string_c in
  let hc = sha1 mar_c in
  let mar_w,variables = iconcat variables in
  let string_w = iutf8 mar_w in
  let w = iS string_w in
  let hw = sha1 mar_w in
  let mar_r,variables = iconcat variables in
  let string_r = iutf8 mar_r in
  let r = iS string_r in
  let hr = sha1 mar_r in
  let store = { 
      vars = {store.vars with c = c; w = w; r = r; };
      hashes = store.hashes;
      macs = store.macs; 
      keys = store.keys; 
      header = store.header} in
  (* Registering new keys *)
  let key_wr,keys = iconcat keys in
  let () = reg_keys (cS store.vars.w) (cS store.vars.r) key_wr in
  let key_cr,_ = iconcat keys in
  let () = reg_keys (cS store.vars.c) (cS store.vars.r) key_cr in
  (* Decrypting variables *)
  let keywr = get_symkey (cS store.vars.w) (cS store.vars.r) in
  let mar_q = unpickle (sym_decrypt keywr encr_q) in
  let string_q = iutf8 mar_q in
  let marred_q = iS string_q in
  let q = int_of_string marred_q in
  let hq = sha1 mar_q in
  let store = { 
      vars = {store.vars with q = q; };
      hashes = {store.hashes with hq = hq; hc = hc; hw = hw; hr = hr; };
      macs = store.macs; 
      keys = {store.keys with key_wr = key_wr; key_cr = key_cr; }; 
      header = store.header} in
  (* Unmarshalling hashes *)
  (* Unmarshalling MACs *)
  let rCommitqcwr,macs = iconcat macs in
  let rFwd,_ = iconcat macs in
  let store = { 
      vars = store.vars;
      hashes = store.hashes;
      macs = {store.macs with rCommitqcwr = rCommitqcwr; rFwd = rFwd; }; 
      keys = store.keys; 
      header = store.header} in
  (* Verification of a MAC from state c_qcwrCommit with variables qcwr*)
  let macheader,maccontent = iconcat store.macs.rCommitqcwr in
  let macsid,ts_mar = iconcat macheader in
  let ts_str = iutf8 ts_mar in
  let ts_string = iS ts_str in
  let ts = int_of_string ts_string in
  let _ = test_inf oldts ts "MAC verification" in
  let oldts = ts in
  let _ = test_eq macsid sid "MAC verification" in
  let mackeycr = get_mackey store.vars.c store.vars.r in
  let content = mac_verify_r_c_qcwrCommit_qcwr ts store mackeycr maccontent in
  (* Verification of a MAC from state w_Fwd_qcwrCommit with variables *)
  let macheader,maccontent = iconcat store.macs.rFwd in
  let macsid,ts_mar = iconcat macheader in
  let ts_str = iutf8 ts_mar in
  let ts_string = iS ts_str in
  let ts = int_of_string ts_string in
  let _ = test_inf oldts ts "MAC verification" in
  let oldts = ts in
  let _ = test_eq macsid sid "MAC verification" in
  let mackeywr = get_mackey store.vars.w store.vars.r in
  let content = mac_verify_r_w_Fwd_qcwrCommit_ ts store mackeywr maccontent in
  let _ = test_eq oldts store.header.ts "Time-stamp verification" in
  (* Verification Ended *)
    let wired = Wired_in_r_start_of_w_Fwd_qcwrCommit___c_qcwrCommit_qcwr__((q,c,w,r),store) in
    wired
  | _ -> failwith "Critical Error"


(* Sending message from c to w *)
(* Message has to be MACed for w r *)
let sendWired_Commit_c_start = fun q -> fun c -> fun w -> fun r -> fun store ->
  let empty_str = cS "" in
  let nil = utf8 empty_str in
  let marred_q = string_of_int q in
  let string_q = cS marred_q in
  let mar_q = utf8 string_q in
  let hq = sha1 mar_q in
  let string_c = cS c in
  let mar_c = utf8 string_c in
  let hc = sha1 mar_c in
  let string_w = cS w in
  let mar_w = utf8 string_w in
  let hw = sha1 mar_w in
  let string_r = cS r in
  let mar_r = utf8 string_r in
  let hr = sha1 mar_r in
  let ts = store.header.ts + 1 in
  let store = { 
      vars = {store.vars with q = q; c = c; w = w; r = r; };
      hashes = {store.hashes with hq = hq; hc = hc; hw = hw; hr = hr; };
      macs = store.macs; 
      keys = store.keys; 
      header = {store.header with ts = ts; }} in
  let ts_string = string_of_int store.header.ts in
  let ts_str = cS ts_string in
  let ts_mar = utf8 ts_str in
  let header = concat store.header.sid ts_mar in
  let keys = nil in
  let keycr = gen_keys (cS store.vars.c) (cS store.vars.r) in
  let keys = concat keycr keys in
  let keycw = gen_keys (cS store.vars.c) (cS store.vars.w) in
  let keys = concat keycw keys in
  (* Generation of a MAC from state c_start to role w *)
  let content = content_c_qcwrCommit_qcwr store.header.ts store in
  let mackeycw = get_mackey store.vars.c store.vars.w in
  let macmsg = mac mackeycw (pickle content) in
  let wCommitqcwr = concat header macmsg in
  (* Generation of a MAC from state c_start to role r *)
  let content = content_c_qcwrCommit_qcwr store.header.ts store in
  let mackeycr = get_mackey store.vars.c store.vars.r in
  let macmsg = mac mackeycr (pickle content) in
  let rCommitqcwr = concat header macmsg in
  let store = { 
      vars = store.vars;
      hashes = store.hashes;
      macs = {store.macs with wCommitqcwr = wCommitqcwr; rCommitqcwr = rCommitqcwr; }; 
      keys = store.keys; 
      header = store.header} in
  let macs = nil in
  let macs = concat store.macs.rCommitqcwr macs in
  let macs = concat store.macs.wCommitqcwr macs in
  let hashes = nil in
  let variables = nil in
  let string_r = cS store.vars.r in
  let mar_r = utf8 string_r in
  let variables = concat mar_r variables in
  let string_w = cS store.vars.w in
  let mar_w = utf8 string_w in
  let variables = concat mar_w variables in
  let string_c = cS store.vars.c in
  let mar_c = utf8 string_c in
  let variables = concat mar_c variables in
  let marred_q = string_of_int store.vars.q in
  let string_q = cS marred_q in
  let mar_q = utf8 string_q in
  let keycw = get_symkey (cS store.vars.c) (cS store.vars.w) in
  let encr_q = sym_encrypt keycw (pickle mar_q) in
  let variables = concat encr_q variables in
  let protocol = concat hashes macs in
  let security = concat keys protocol in
  let payload = concat variables security in
  let visib = cS "c_qcwrCommit_qcwr__" in
  let visib_string = utf8 visib in
  let content = concat visib_string payload in
  let msg = base64 (concat header content) in
  let () = psend  store.vars.w msg in 
  store 

(* Sending message from w to r *)
(* Message has to be MACed for r *)
let sendWired_Fwd_w_qcwrCommit = fun store ->
  let empty_str = cS "" in
  let nil = utf8 empty_str in
  let ts = store.header.ts + 1 in
  let store = { 
      vars = store.vars;
      hashes = store.hashes;
      macs = store.macs; 
      keys = store.keys; 
      header = {store.header with ts = ts; }} in
  let ts_string = string_of_int store.header.ts in
  let ts_str = cS ts_string in
  let ts_mar = utf8 ts_str in
  let header = concat store.header.sid ts_mar in
  let keys = nil in
  let keys = concat store.keys.key_cr keys in
  let keywr = gen_keys (cS store.vars.w) (cS store.vars.r) in
  let keys = concat keywr keys in
  (* Generation of a MAC from state w_qcwrCommit to role r *)
  let content = content_w_Fwd_qcwrCommit_ store.header.ts store in
  let mackeywr = get_mackey store.vars.w store.vars.r in
  let macmsg = mac mackeywr (pickle content) in
  let rFwd = concat header macmsg in
  let store = { 
      vars = store.vars;
      hashes = store.hashes;
      macs = {store.macs with rFwd = rFwd; }; 
      keys = store.keys; 
      header = store.header} in
  let macs = nil in
  let macs = concat store.macs.rFwd macs in
  let macs = concat store.macs.rCommitqcwr macs in
  let hashes = nil in
  let variables = nil in
  let string_r = cS store.vars.r in
  let mar_r = utf8 string_r in
  let variables = concat mar_r variables in
  let string_w = cS store.vars.w in
  let mar_w = utf8 string_w in
  let variables = concat mar_w variables in
  let string_c = cS store.vars.c in
  let mar_c = utf8 string_c in
  let variables = concat mar_c variables in
  let marred_q = string_of_int store.vars.q in
  let string_q = cS marred_q in
  let mar_q = utf8 string_q in
  let keywr = get_symkey (cS store.vars.w) (cS store.vars.r) in
  let encr_q = sym_encrypt keywr (pickle mar_q) in
  let variables = concat encr_q variables in
  let protocol = concat hashes macs in
  let security = concat keys protocol in
  let payload = concat variables security in
  let visib = cS "w_Fwd_qcwrCommit___c_qcwrCommit_qcwr__" in
  let visib_string = utf8 visib in
  let content = concat visib_string payload in
  let msg = base64 (concat header content) in
  let () = psend  store.vars.r msg in 
  store