import num
import list

-- Implementation of RSA encryption algorithm.
-- Reference: https://simple.wikipedia.org/wiki/RSA_(algorithm)


-- To use with randomly generated prime numbers, call `randKeys`
-- with a number, representing how many digits the generated primes 
-- will have, and a pseudorandom number generator, created 
-- by calling `seed` with a natural number. 

-- To manually input prime numbers, call `getKeys` 
-- with two prime numbers. 

-- Output: Both of the above functions output two pairs of two 
-- natural numbers. The first pair is the public key, the 
-- second is the private key. These keys, along with the `encrypt` 
-- and `decrypt` functions can be used to encrypt and decrypt 
-- lists of natural numbers.

encrypt : N * N -> List(N) -> List(N)
encrypt key xs = each (encrypt1 key, xs)

decrypt : N * N -> List(N) -> List(N)
decrypt = encrypt


randKeys : N -> Gen -> (N * N) * (N * N)
randKeys 0 g = randKeys 1 g
randKeys d g = 
   let   p1 = genPrime (10^(d.-1),10^d) g,
         p2 = genPrime (10^(d.-1),10^d) (snd p1)
   in getKeys (fst p1) (fst p2) 

genPrime : (N * N) -> Gen -> (N * Gen)
genPrime r g = let p = random(r,g) in 
   {? p                          if isPrime (fst p)
   ,  genPrime r (snd p)         otherwise
   ?}

fst : (a * b) -> a 
fst (a,_) = a 

snd : (a * b) -> b 
snd (_,b) = b 


-- takes two primes, returns a pair of pairs containing the RSA public/private keys
-- prime -> prime -> (public key, private key)

getKeys : N -> N -> (N * N) * (N * N)
getKeys p1 p2 =
  let m = p1 * p2,
      totient = (p1 .- 1)*(p2 .- 1),
      e = getPubExp 2 totient
  in ((m, e), (m, getPrivExp e totient))

-- guess -> totient -> e
getPubExp : N -> N -> N
getPubExp e totient =
  {? e                          if gcd(e, totient) == 1
   , getPubExp (e+1) totient    otherwise
  ?}

gcd : N*N -> N
gcd (a, 0) = a
gcd (a, b) = gcd (b, a mod b)

getPrivExp : N -> N -> N
getPrivExp e totient =
  let t = inverse (0,1) (totient,e)
  in {? abs t              if t>=0
      , abs (t+totient)    otherwise
     ?}


-- Implemented using Extended Euclidean Algorithm (reference:
-- https://en.wikipedia.org/wiki/Extended_Euclidean_algorithm#Computing_multiplicative_inverses_in_modular_structures)
inverse : (Z * Z) -> (Z * Z) -> Z
inverse (t,newt) (r,newr) =
  {? t                                                                    if newr==0
   , let q = r // newr in (inverse (newt, t-q*newt) (newr,r-q*newr))      otherwise
  ?}

-- encrypt1 : msg -> public key (mod,exp) -> encrypted msg
-- encrypts one single number
encrypt1 : Nat * Nat -> Nat -> Nat
encrypt1 (m, e) msg = modPower msg e m

-- decrypts one single number
decrypt1 : Nat * Nat -> Nat -> Nat
decrypt1 = encrypt1

-- modPower : n -> power -> modulus -> nat
-- Exponentiating by squaring algorithm reference:
-- https://simple.wikipedia.org/wiki/Exponentiation_by_squaring
modPower : Nat -> Nat -> Nat -> Nat
modPower n p m =
  {? 1                                     if p==0
   , n % m                                 if p==1
   , (modPower (n^2) (p//2) m) % m         if (even p)
   , (n * (modPower (n^2) (p//2) m)) % m   if (p>2) && (odd p)
  ?}