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_qcwCommit),({q,c,w}Commit{c,w}),(w_start,w_qcwCommit))(w_qcwCommit,w_Ack_qcwCommit),({}Ack{}),(c_qcwCommit,c_Ack_qcwCommit))(c_Ack_qcwCommit,c_Reveal_qcwAck),({}Reveal{q}),(w_Ack_qcwCommit,w_Reveal_qcwAck))"))


type vars = {
  c : principal ;
  w : principal ;
  q : int }
type hashes = {
  hc : bytes ;
  hw : bytes ;
  hq : bytes }
type macs = {
  wCommitqcw : bytes ;
  cAck : bytes ;
  wReveal : bytes }
type keys = {
  key_cw : bytes ;
  key_wc : 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 = "" ; q = 0 };
   hashes = { hc = utf8 (cS "") ; hw = utf8 (cS "") ; hq = utf8 (cS "") };
   macs = { wCommitqcw = utf8 (cS "") ; cAck = utf8 (cS "") ; wReveal = utf8 (cS "") };
   keys = { key_cw = utf8 (cS "") ;
  key_wc = 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 ; q = 0 };
   hashes = { hc = utf8 (cS "") ; hw = utf8 (cS "") ; hq = utf8 (cS "") };
   macs = { wCommitqcw = utf8 (cS "") ; cAck = utf8 (cS "") ; wReveal = utf8 (cS "") };
   keys = { key_cw = utf8 (cS "") ;
  key_wc = 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)
| Send_Ack of (principal * bytes * int)
| Send_Reveal of (principal * bytes * int)
| Recv_Commit of (principal * bytes * int * principal * principal)
| Recv_Ack of (principal * bytes * int)
| Recv_Reveal of (principal * bytes * int * int)



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

let content_c_qcwCommit_qcw = fun ts store ->
  let empty_str = cS "" in
  let nil = utf8 empty_str in
  let hashes = nil 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_qcwCommit" 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_w_Ack_qcwCommit_ = fun ts store ->
  let empty_str = cS "" in
  let nil = utf8 empty_str in
  let hashes = nil in
  let nextstate = cS "w_Ack_qcwCommit" 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_Reveal_qcwAck_ = fun ts store ->
  let empty_str = cS "" in
  let nil = utf8 empty_str in
  let hashes = nil in
  let nextstate = cS "c_Reveal_qcwAck" 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_qcwCommit_qcw = fun ts store k m ->
  let content = content_c_qcwCommit_qcw 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_c_w_Ack_qcwCommit_ = fun ts store k m ->
  let content = content_w_Ack_qcwCommit_ 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_w_c_Reveal_qcwAck_ = fun ts store k m ->
  let content = content_c_Reveal_qcwAck_ 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_qcwCommit_qcw__ of (principal * principal) * store

type wired_c_qcwCommit =
  | Wired_in_c_qcwCommit_of_w_Ack_qcwCommit___ of (unit) * store

type wired_w_Ack_qcwCommit =
  | Wired_in_w_Ack_qcwCommit_of_c_Reveal_qcwAck___ of (int) * 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_qcwCommit_qcw__" ->
  let variables,security = iconcat payload in
  let keys,protocol = iconcat security in
  let hashes,macs = iconcat protocol in
  (* Unmarshalling variables *)
  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 store = { 
      vars = {store.vars with c = c; w = w; };
      hashes = store.hashes;
      macs = store.macs; 
      keys = store.keys; 
      header = store.header} in
  (* Registering new keys *)
  let key_cw,_ = iconcat keys in
  let () = reg_keys (cS store.vars.c) (cS store.vars.w) key_cw in
  (* Decrypting variables *)
  let store = { 
      vars = store.vars;
      hashes = {store.hashes with hc = hc; hw = hw; };
      macs = store.macs; 
      keys = {store.keys with key_cw = key_cw; }; 
      header = store.header} in
  (* Unmarshalling hashes *)
  let hq,_ = iconcat hashes in
  (* Unmarshalling MACs *)
  let wCommitqcw,_ = iconcat macs in
  let store = { 
      vars = store.vars;
      hashes = {store.hashes with hq = hq; };
      macs = {store.macs with wCommitqcw = wCommitqcw; }; 
      keys = store.keys; 
      header = store.header} in
  (* Verification of a MAC from state c_qcwCommit with variables qcw*)
  let macheader,maccontent = iconcat store.macs.wCommitqcw 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_qcwCommit_qcw 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_qcwCommit_qcw__((c,w),store) in
    wired

let receiveWired_c_qcwCommit : store -> wired_c_qcwCommit = fun store ->
  let empty_str = cS "" in
  let nil = utf8 empty_str in
  let msg = precv  store.vars.c 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 oldts = store.header.ts in
  let _ = test_inf oldts ts "Replay attack!" in
  let _ = test_eq store.header.sid sid "Session confusion attack!" in
  let store = { 
      vars = store.vars;
      hashes = store.hashes;
      macs = store.macs; 
      keys = store.keys; 
      header = {store.header with ts = ts; }} in
  let tag,payload = iconcat content in
  match iS (iutf8 tag) with
  | "w_Ack_qcwCommit___" ->
  let variables,security = iconcat payload in
  let keys,protocol = iconcat security in
  let hashes,macs = iconcat protocol in
  (* Unmarshalling variables *)
  (* Registering new keys *)
  let key_wc,_ = iconcat keys in
  let () = reg_keys (cS store.vars.w) (cS store.vars.c) key_wc in
  (* Decrypting variables *)
  let store = { 
      vars = store.vars;
      hashes = store.hashes;
      macs = store.macs; 
      keys = {store.keys with key_wc = key_wc; }; 
      header = store.header} in
  (* Unmarshalling hashes *)
  (* Unmarshalling MACs *)
  let cAck,_ = iconcat macs in
  let store = { 
      vars = store.vars;
      hashes = store.hashes;
      macs = {store.macs with cAck = cAck; }; 
      keys = store.keys; 
      header = store.header} in
  (* Verification of a MAC from state w_Ack_qcwCommit with variables *)
  let macheader,maccontent = iconcat store.macs.cAck 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 store.header.sid "MAC verification" in
  let mackeywc = get_mackey store.vars.w store.vars.c in
  let content = mac_verify_c_w_Ack_qcwCommit_ ts store mackeywc maccontent in
  let _ = test_eq oldts store.header.ts "Time-stamp verification" in
  (* Verification Ended *)
    let wired = Wired_in_c_qcwCommit_of_w_Ack_qcwCommit___((),store) in
    wired

let receiveWired_w_Ack_qcwCommit : store -> wired_w_Ack_qcwCommit = 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 oldts = store.header.ts in
  let _ = test_inf oldts ts "Replay attack!" in
  let _ = test_eq store.header.sid sid "Session confusion attack!" in
  let store = { 
      vars = store.vars;
      hashes = store.hashes;
      macs = store.macs; 
      keys = store.keys; 
      header = {store.header with ts = ts; }} in
  let tag,payload = iconcat content in
  match iS (iutf8 tag) with
  | "c_Reveal_qcwAck___" ->
  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
  (* Registering new keys *)
  (* Decrypting variables *)
  let mar_q = unpickle (sym_decrypt (get_symkey (cS store.vars.c) (cS store.vars.w)) 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 marred_q = string_of_int q in
  let string_q = cS marred_q in
  let mar_q = utf8 string_q in
    let () = sha1_verify mar_q store.hashes.hq in
  let store = { 
      vars = {store.vars with q = q; };
      hashes = store.hashes;
      macs = store.macs; 
      keys = store.keys; 
      header = store.header} in
  (* Unmarshalling hashes *)
  (* Unmarshalling MACs *)
  let wReveal,_ = iconcat macs in
  let store = { 
      vars = store.vars;
      hashes = store.hashes;
      macs = {store.macs with wReveal = wReveal; }; 
      keys = store.keys; 
      header = store.header} in
  (* Verification of a MAC from state c_Reveal_qcwAck with variables *)
  let macheader,maccontent = iconcat store.macs.wReveal 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 store.header.sid "MAC verification" in
  let mackeycw = get_mackey store.vars.c store.vars.w in
  let content = mac_verify_w_c_Reveal_qcwAck_ ts store mackeycw maccontent in
  let _ = test_eq oldts store.header.ts "Time-stamp verification" in
  (* Verification Ended *)
    let wired = Wired_in_w_Ack_qcwCommit_of_c_Reveal_qcwAck___((q),store) in
    wired
  | _ -> failwith "Critical Error"


(* Sending message from c to w *)
(* Message has to be MACed for w *)
let sendWired_Commit_c_start = fun q -> fun c -> fun w -> 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 ts = store.header.ts + 1 in
  let store = { 
      vars = {store.vars with q = q; c = c; w = w; };
      hashes = {store.hashes with hq = hq; hc = hc; hw = hw; };
      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 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_qcwCommit_qcw store.header.ts store in
  let mackeycw = get_mackey store.vars.c store.vars.w in
  let macmsg = mac mackeycw (pickle content) in
  let wCommitqcw = concat header macmsg in
  let store = { 
      vars = store.vars;
      hashes = store.hashes;
      macs = {store.macs with wCommitqcw = wCommitqcw; }; 
      keys = store.keys; 
      header = store.header} in
  let macs = nil in
  let macs = concat store.macs.wCommitqcw macs in
  let hashes = nil in
  let hashes = concat store.hashes.hq hashes in
  let variables = nil 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 protocol = concat hashes macs in
  let security = concat keys protocol in
  let payload = concat variables security in
  let visib = cS "c_qcwCommit_qcw__" 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 c *)
(* Message has to be MACed for c *)
let sendWired_Ack_w_qcwCommit = 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 keywc = gen_keys (cS store.vars.w) (cS store.vars.c) in
  let keys = concat keywc keys in
  (* Generation of a MAC from state w_qcwCommit to role c *)
  let content = content_w_Ack_qcwCommit_ store.header.ts store in
  let mackeywc = get_mackey store.vars.w store.vars.c in
  let macmsg = mac mackeywc (pickle content) in
  let cAck = concat header macmsg in
  let store = { 
      vars = store.vars;
      hashes = store.hashes;
      macs = {store.macs with cAck = cAck; }; 
      keys = store.keys; 
      header = store.header} in
  let macs = nil in
  let macs = concat store.macs.cAck macs in
  let hashes = nil in
  let variables = nil in
  let protocol = concat hashes macs in
  let security = concat keys protocol in
  let payload = concat variables security in
  let visib = cS "w_Ack_qcwCommit___" 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.c msg in 
  store 

(* Sending message from c to w *)
(* Message has to be MACed for w *)
let sendWired_Reveal_c_Ack_qcwCommit = 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
  (* Generation of a MAC from state c_Ack_qcwCommit to role w *)
  let content = content_c_Reveal_qcwAck_ store.header.ts store in
  let mackeycw = get_mackey store.vars.c store.vars.w in
  let macmsg = mac mackeycw (pickle content) in
  let wReveal = concat header macmsg in
  let store = { 
      vars = store.vars;
      hashes = store.hashes;
      macs = {store.macs with wReveal = wReveal; }; 
      keys = store.keys; 
      header = store.header} in
  let macs = nil in
  let macs = concat store.macs.wReveal macs in
  let hashes = nil in
  let variables = nil 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_Reveal_qcwAck___" 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