(* Options to decide which script to generate *)

let detailed_lockbox = ref false
let corrected = ref false
let signhash_without_enc = ref false

let rec succ n s =
  if n = 0 then print_string s else
  begin
    print_string "succ(";
    succ (n-1) s;
    print_string ")"
  end

let revoke n = 
  print_string "

      (* Next, we move from version "; print_int n; print_string " to version "; print_int (n+1); print_string ". *) 
      | 
       (
         (* receive a revoke request for group g; 
            newreaders and newwriters and the new lists of allowed readers and writers *)
         in(net, (=revoke,=g,newreaders,newwriters));       

         (* read state for version "; print_int n; print_string " *)
         in(currentstate, (=";
  succ n "zero";
  print_string ",oldreaders,oldwriters,oldlk,oldsk));

         (* wind old lk to new lk *)
         let newlk = exp(oldlk,ownerprivkey) in              

         (* generate new sk *)
         new seed3; let newsk = (d(seed3,newlk),N(seed3)) in        

         (* assert the state at version "; print_int (n+1); print_string " *)
         event isstate(g,";
  succ (n+1) "zero";
  print_string ",newreaders,newwriters);

         (* store state for version "; print_int (n+1); print_string " on channel currentstate *)
         out(currentstate, (";
  succ (n+1) "zero";
  print_string ",newreaders,newwriters,newlk,newsk))    
       )"

let rec revokeTill n = 
  if n = 0 then () else 
    begin
      revokeTill (n-1);
      revoke (n-1)
    end

let rec repeat k s = 
  if k = 0 then () else 
    begin
      print_string s;
      repeat (k-1) s
    end

let rec unwind n k s = 
  print_string "         if v = "; succ k "zero"; print_string " then (\n";
  repeat (n-k) 
"            (* unwind lk *)
            let lk = exp(lk,ownerpubkey) in\n\n";
  print_string s;
  if k < n then 
    begin
      print_string "         )\n        |\n";
      unwind n (k+1) s
    end
  else
    print_string "         )\n"

let rec unwindTill n s = 
  if n > 0 then 
    begin
      print_string "      if vx = "; succ n "zero"; print_string " then (\n";
      unwind n 0 s;
      print_string "      )\n      |\n";
      unwindTill (n-1) s
    end
  else
    begin
      print_string "      if vx = zero then (
         if v = zero then (\n";
      print_string s;
      print_string "      ))"
    end

let plutus n =
  print_string "(* Generated script of the protocol Plutus with ";
  print_int n;
  print_string " versions";
  if !corrected then
    print_string ", corrected"
  else
    print_string ", original";
  if !detailed_lockbox then
    print_string ", detailed lockbox";
  if !signhash_without_enc then
    print_string ", contents hashed and signed unencrypted";
  print_string " *)

param redundantHypElim = beginOnly.
param traceBacktracking = false.

(** FUNCTION SYMBOLS **)

(* exponents *)
fun e/2.
fun d/2.

(* modulus; a random RSA seed (p,q) is modeled as a fresh name s *)
fun N/1.

(* exponentiation modulo *)
fun exp/2.

(* symmetric encryption *)
fun enc/2.

(* hash function *)
fun hash/1.

(** ALGORITHMS AND EQUATIONS *)

(* derive seed given RSA exponents and modulus *)
reduc crack(e(x,y),d(x,y),N(x)) = x.                             

(* verify RSA signature, etc. *)
equation exp(exp(z,(d(x,y),N(x))),(e(x,y),N(x))) = z.            
equation exp(exp(z,(e(x,y),N(x))),(d(x,y),N(x))) = z.            

(* derive exponent from modulus and base *)
reduc genExp(N(x),y) = e(x,y).                                    

(* decrypt symmetric encryption *)
reduc dec(enc(x,y),y) = x.                                        

(** VERSION NUMBERS *)

data zero/0.
data succ/1.

pred geq/2.
clauses
        geq:x,x;
	geq:x,y -> geq:succ(x),y.

(** ACCESS CONTROL LISTS *)

data cons/2.
data nil/0.

pred member/2.
clauses
  member:x,cons(x,y);
  member:x,y -> member:x,cons(z,y).

(** CHANNEL DECLARATIONS *)

(* The network and file system are under control of the adversary. *)
free net,fs.

(* The adversary can choose actions. *)
free newgroup.
free revoke.
free rkeyreq, wkeyreq.
free corrupt. 

(* Keys are communicated securely to users. *)
private fun rprivchannel/1.
private fun wprivchannel/1.

(** QUERIES *)

(* rx is a reader of group gx at version vx only if it belongs to the readers rlx of gx at vx. *) 
query ev:isreader(rx,gx,vx) ==>
                  ev:isgroup(gx)   
                & ev:isstate(gx,vx,rlx,wlx)
	        & member:rx,rlx.

(* wx is a writer of group gx at version vx only if it belongs to the writers wlx of gx at vx. *) 
query ev:iswriter(wx,gx,vx) ==>
                  ev:isgroup(gx)   
                & ev:isstate(gx,vx,rlx,wlx)
  	        & member:wx,wlx.

(* See definition of secrecy. *)
query let mx = m[g = gx; v = vy]; 
      attacker:mx ==> 
		  ev:corrupt(rx,gx,vx)
		& geq:vx,vy
      	      	& ev:isreader(rx,gx,vx)  
              |
		  ev:corrupt(rx,gx,vx)
		& geq:vx,vy
      	      	& ev:iswriter(rx,gx,vx).

(* See definition of integrity. *)
query ev:gets(rx,mx,gx,vx) ==> 
                  ev:iswriter(wx,gx,vx) 
                & ev:puts(wx,mx,gx,vx) 
              |
                  ev:iswriter(wx,gx,vx)
                & ev:corrupt(wx,gx,vx).

(** CODE FOR OWNERS *)

let processOwr =
    (* create owner's RSA key pair *)
    new seed1; new seed2;                                   
    let ownerpubkey = (e(seed1,seed2),N(seed1)) in                
    let ownerprivkey = (d(seed1,seed2),N(seed1)) in

    (* publish owner's RSA public key *)
    out(net,ownerpubkey);

    (
    ! (* receive a new group creation request;
         initreaders and initwriters are the initial lists of allowed readers and writers, respectively *)
      in(net, (=newgroup,initreaders,initwriters));         

      (* create a new group g *)
      new g;

      (* assert that g is a group *)
      event isgroup(g);

      (* publish the name g *)
      out(net,g);

      (* create a private channel for the current state for group g *)
      new currentstate;

      (
       ( (* create initial lk *)
	 new initlk;       
                                     
	 (* generate initial sk *)
         new seed3; let initsk = (d(seed3,initlk),N(seed3)) in        
   
	 (* assert the state at version 0 *)
         event isstate(g,zero,initreaders,initwriters);

	 (* store state for version 0 on channel currentstate *)
         out(currentstate, (zero,initreaders,initwriters,initlk,initsk))
       )";
    revokeTill n;
    print_string "
      |
       (
       ! (* receive read key request for reader r and group g *)
         in(net, (=rkeyreq,r,=g));                            

         (* get the current state *)    
         in(currentstate, (v,readers,writers,lk,sk));
         out(currentstate, (v,readers,writers,lk,sk));

         (* check that the reader r is allowed *)
         if member:r,readers then (                             

            (* assert that r is a reader for group g and version v *)
            event isreader(r,g,v); 

            (* send lk and owner's public key to r *)
            out(rprivchannel(r), (g,v,lk,ownerpubkey)) )
       )
      |
       ( 
       ! (* receive write key request for writer w and group g *)
         in(net, (=wkeyreq,w,=g));                             

         (* get the current state *)   
         in(currentstate, (v,readers,writers,lk,sk));
         out(currentstate, (v,readers,writers,lk,sk));

         (* check that the writer w is allowed *)
         if member:w,writers then (                             
                                                               
             (* sign the modulus *)  
             let (dx,n) = sk in";
    if !corrected then
      print_string " let sn = exp(hash((n,g,v)),ownerprivkey) in"
    else
      print_string " let sn = exp(hash(n),ownerprivkey) in";
    print_string "

             (* assert that w is a writer for group g and version v *)
	     event iswriter(w,g,v);

             (* send lk, sk, and signed modulus to w *)
             out(wprivchannel(w), (g,v,lk,sk,sn)) )
       )
      )
    ).

(** CODE FOR WRITERS *)

let processGoodWtr =
    (* initiate a writer w for group g *)
    in(net, (w,g)); 

    (* send write key request *)
    out(net, (wkeyreq,w,g));                                      

    (* obtain lk, sk, and signed modulus *)
    in(wprivchannel(w), (=g,v,lk,sk,sn));                   

    (* create data to write *)
    new m;

    (* encrypt *)";
  if !detailed_lockbox then 
    print_string "
    new k;
    let encx = (enc(k,lk),enc(m,k)) in

    (* sign *)"
  else
    print_string "
    let encx = enc(m,lk) in

    (* sign *)";
  if !signhash_without_enc then print_string "
    let sencx = exp(hash(m),sk) in"
  else print_string "
    let sencx = exp(hash(encx),sk) in";
  print_string "

    (* assert that data m has been written by w for group g at version v *)
    event puts(w,m,g,v);

    (* send content to the file system *)
    let (dx,n) = sk in out(fs, (g,v,n,sn,encx,sencx)).

let processBadWtr = 
    (* initiate a writer w for group g *)
    in(net, (w,g)); 

    (* send write key request *)
    out(net, (wkeyreq,w,g));                                      

    (* obtain lk, sk, and signed modulus *)
    in(wprivchannel(w), (=g,v,lk,sk,sn));                   

    (* receive corrupt request for w *)
    in(net, =(corrupt,w));             

    (* assert that w has been corrupted for group g at version v *)                           
    event corrupt(w,g,v);

    (* leak lk, sk, and signed modulus *)
    out(net, (lk,sk,sn)).                                   

(** CODE FOR READERS *)

let processGoodRdr = 
    (* initiate a reader r for group g *)
    in(net, (r,g));

    (* send read key request *)
    out(net, (rkeyreq,r,g));   

    (* obtain lk and owner's public key *)
    in(rprivchannel(r), (=g,vx,lk,ownerpubkey));      

    (* obtain header and content from file system *)
    in(fs, (=g,v,n,sn,encx,sencx));                               

    (* verify signature in header *)";
  if !corrected then
    print_string "
    if hash((n,g,v)) = exp(sn,ownerpubkey) then"
  else
    print_string "
    if hash(n) = exp(sn,ownerpubkey) then";
  print_string "

    (
";
    unwindTill n
	(if ((!signhash_without_enc) && (!detailed_lockbox)) then
"            (* derive vk *) 
            let vk = (genExp(n,lk),n) in

            (* decrypt to obtain data *)
            let (encl,encf) = encx in
            let k = dec(encl,lk) in
	    let x = dec(encf,k) in

            (* verify signature of data *)
            if hash(x) = exp(sencx,vk) then

            (* assert that reader r read data x for group g and version v *)
            event gets(r,x,g,v)\n"
	 else if ((!signhash_without_enc) && (not (!detailed_lockbox))) then
"            (* derive vk *)            
            let vk = (genExp(n,lk),n) in

            (* decrypt to obtain data *)
	    let x = dec(encx,lk) in

            (* verify signature of data *)
            if hash(x) = exp(sencx,vk) then

            (* assert that reader r read data x for group g and version v *)
            event gets(r,x,g,v)\n"
	 else if (!detailed_lockbox) then
"            (* derive vk *)                       
            let vk = (genExp(n,lk),n) in

            (* verify signature of encryption *)
            if hash(encx) = exp(sencx,vk) then

            (* decrypt to obtain data *)
            let (encl,encf) = encx in
            let k = dec(encl,lk) in
	    let x = dec(encf,k) in

            (* assert that reader r read data x for group g and version v *)
            event gets(r,x,g,v)\n"
	    else
"            (* derive vk *)                       
            let vk = (genExp(n,lk),n) in

            (* verify signature of encryption *)
            if hash(encx) = exp(sencx,vk) then

            (* decrypt to obtain data *)
            let x = dec(encx,lk) in

            (* assert that reader r read data x for group g and version v *)
            event gets(r,x,g,v)\n");
    print_string "
    ).

let processBadRdr =
    (* initiate a reader r for group g *)
    in(net, (r,g));

    (* send read key request *)
    out(net, (rkeyreq,r,g));   

    (* obtain lk and owner's public key *)
    in(rprivchannel(r), (=g,v,lk,ownerpubkey));      

    (* receive corrupt request for r *)
    in(net, =(corrupt,r));                                       

    (* assert that r has been corrupted for group g at version v *)
    event corrupt(r,g,v);

    (* leak lk *)
    out(net,lk).                                                

(** TOP LEVEL PROCESS *)

process
         ( processOwr
	 | (! processGoodWtr) | (! processBadWtr)
         | (! processGoodRdr) | (! processBadRdr)  
         ) 

"

let _ =
  Arg.parse 
    [ "-corrected", Arg.Set corrected, "\t\tcorrected script, with modulus, version, and group signed";
      "-signhashwithoutenc", Arg.Set signhash_without_enc, "\t\tcontents are hashed and signed unencrypted";
      "-detailedlockbox", Arg.Set detailed_lockbox, "\t\tthe lockbox key encrypts a fresh key, not the file itself"]
 (fun s -> 
  try 
    plutus (int_of_string s)
  with _ ->
    print_string "Error. Usage: plutusgen [options] n\nwhere n is an integer\n")
  "Plutus script generator"