(*
  This file contains the fundamental cryptography operations
*)

#load "nums.cma";;

exception CryptoException

open Big_int;;

(* define function shorthands *)
let bsucc = succ_big_int
let bpred = pred_big_int
let badd = add_big_int
let bsub = sub_big_int
let bmult = mult_big_int
let bdiv = div_big_int
let beq = eq_big_int
(* let bsq = square_big_int *)
let bsq = fun x -> bmult x x
let ( % ) = mod_big_int
let bquomod = quomod_big_int
let bsign = sign_big_int
let bge = ge_big_int
let bgt = gt_big_int
let string2b = big_int_of_string

(* convenience printing functions *)
let prs x = print_string ("Big_int: " ^ (string_of_big_int x)); print_newline ();;
let prs4 (x1,x2,x3,x4) = 
	print_string ("B1: " ^ (string_of_big_int x1));
	print_newline ();
	print_string ("B2: " ^ (string_of_big_int x2));
	print_newline ();
	print_string ("B3: " ^ (string_of_big_int x3));
	print_newline ();
	print_string ("B4: " ^ (string_of_big_int x4));
	print_newline ();;
let ps x = print_string x; print_newline();;

(* short hands for 0, 1, and 2 *)
let b0 = zero_big_int
let b1 = unit_big_int
let b2 = bsucc b1
let b3 = bsucc b2
let b4 = bsucc b3

(* staged mult *)
let bmult' x y = .< bmult x (.~ y) >.;;

(* another staged mult *)
let bmult'' x y = .< bmult (.~ x) (.~ y) >.;;

(* convert x to positive module m *)
let pos x m =
	if bsign x = -1 then badd x m
			else x

(* even function for big int *)
let beven x = beq (x % b2) b0

let even x = (x mod 2) = 0

(* staged square function *)
let bsq' x = .< let y = .~ x in bsq y >.

(* power function for big int *)
let rec bpow x y =
	if beq y b1 then x
		    else if beven y then bsq (bpow x (bdiv y b2))
				    else bmult x (bpow x (bpred y));;

(* power function for big int to int *)
let rec bpowint x y =
	if y=1 then x
	       else if even y then bsq (bpowint x (y/2))
			      else bmult x (bpowint x (y-1));;

(* power-mod function *)
let rec bpowmod x y z =
	if beq y b1 then x
		    else let (q, r) = bquomod y b2 in
			 if beq r b0 then let w = bpowmod x q z in 
					  (bsq w) % z
				     else (bmult x (bpowmod x (bpred y) z)) % z

(* staged version of mod function *)
let ( %% ) = fun x y -> .< (.~ x) % y >.

(* staged version of power-mod function *) 
let rec bpowmod' x y z =
	if beq y b1 then x
		    else let (q, r) = bquomod y b2 in
		    	 if beq r b0 then .< let w = .~(bpowmod' x q z) in
			 		     (bsq w) % z >.
				     else .< let x' = .~ x in
				             let w = .~(bpowmod' .<x'>. (bpred y) z) in
				             (bmult x' w) % z >.

(* convert a big int into a big-endian boolean list representation *)
let rblit x =
	let l = ref [] in
	let v = ref x in
	while bgt !v b0 do
		let (q,r) = bquomod !v b2 in
		begin
			if beq r b1 then l := true::(!l)
				    else l := false::(!l);
			v := q
		end
	done;
	!l

(* convert a big int into a small-endian boolean list representation *)
let blit x = List.rev (rblit x)

(* helper function for bpowmodAdd *)
(* square p and multiply to x if head of list is true *)
let rec sqmult x l p z =
	match l with
	  [] -> x
	| h::t -> if h then sqmult ((bmult x p) % z) t ((bsq p) % z) z
		       else sqmult x t ((bsq p) % z) z

(* a better power-mod function *)
let bpowmodAdd x y z =
	let rec sqmult x l p =
		match l with
		  [] -> x
		| h::t -> if h then sqmult ((bmult x p) % z) t ((bsq p) % z)
			       else sqmult x t ((bsq p) % z)
	in
	let l = blit y in
	sqmult b1 l x

(* a staged better power-mod function *)
let bpowmodAdd' x y z =
	let rec sqmult' x l p =
		match l with
		  [] -> x
		| h::t -> if h then
		     .< let y = .~ p in
			.~(sqmult' .<(bmult (.~ x) y) % z>. t .<(bmult y y) % z>.) >.
		    else sqmult' x t .<(bsq .~ p) % z>.
	in
	let l = blit y in
	sqmult' .<b1>. l x

(* extended Euclid's algorithm for big int *)
let rec bexteuclid a b =
	if beq b b0 then (b1, b0)
		    else let (q, r) = bquomod a b in
			 let (x1, x2) = bexteuclid b r in
			 (x2, bsub x1 (bmult q x2))

(* inverse of p mod m *)
let inv p m =
	let x = fst (bexteuclid p m) in
	if bsign x = -1 then badd x m
			else x

(* Prepare numbers for CRT use *)
let bigypq (x1,x2) n p q = ((bmult x1 p) % n, (bmult x2 q) % n)

(* big int version of decryption with CRT *)
let bdecCRT c d p q n =
	let d1 = d % (bpred p) in
	let d2 = d % (bpred q) in
	let x = bexteuclid p q in
	let (y1, y2) = bigypq x n p q in
	let mp = bpowmod c d1 p in
	let mq = bpowmod c d2 q in
		(badd (bmult mp y2) (bmult mq y1)) % n;;

(* big int staged version of decryption with CRT *)
let bdecCRT' c d p q n =
	let d1 = d % (bpred p) in
	let d2 = d % (bpred q) in
	let x = bexteuclid p q in
	let (y1, y2) = bigypq x n p q in
	let mp = bpowmod' c d1 p in
	let mq = bpowmod' c d2 q in
    		.< (badd (bmult (.~ mp) y2) (bmult (.~ mq) y1)) % n >.;;

(* big int version of decryption with CRT *)
let bdecCRTAdd c d p q n =
	let d1 = d % (bpred p) in
	let d2 = d % (bpred q) in
	let x = bexteuclid p q in
	let (y1, y2) = bigypq x n p q in
	let mp = bpowmodAdd c d1 p in
	let mq = bpowmodAdd c d2 q in
		(badd (bmult mp y2) (bmult mq y1)) % n;;

(* big int staged version of decryption with CRT *)
let bdecCRTAdd' c d p q n =
	let d1 = d % (bpred p) in
	let d2 = d % (bpred q) in
	let x = bexteuclid p q in
	let (y1, y2) = bigypq x n p q in
	let mp = bpowmodAdd' c d1 p in
	let mq = bpowmodAdd' c d2 q in
    		.< (badd (bmult (.~ mp) y2) (bmult (.~ mq) y1)) % n >.;;