import num
import list

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

-- To use, first call `getKeys` with two prime numbers, which returns
-- two pairs. 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

-- 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)
  ?}