#!/usr/bin/runghc

import CryptoMath
import Data.Char
import Data.Maybe (fromJust)
import RabinMiller
import System.Random

-- Note: The keysize here is 2^64, so you can see the numbers on the screen. 
-- Normal key size is 2^4096...

genLargePrime :: IO Integer
genLargePrime = do
  bigRand <- randomRIO (2 ^ 63, 2 ^ 64) :: IO Integer
  if isPrime bigRand then return bigRand else genLargePrime

genPQ :: IO (Integer, Integer)
genPQ = do
  p <- genLargePrime
  q <- genLargePrime
  if p /= q then return (p, q) else genPQ

genCoPrime :: Integer -> IO Integer
genCoPrime n = do
  bigRand <- randomRIO (2 ^ 63, 2 ^ 64) :: IO Integer
  if gcd' bigRand n == 1 then return bigRand else genCoPrime n

genKeys :: IO ((Integer, Integer), (Integer, Integer))
genKeys = do
  (p, q) <- genPQ
  let n = p * q
      phiN = (p - 1) * (q - 1)
  e <- genCoPrime phiN
  let d = fromJust (modInv e phiN)
  return ((n, e), (n, d))

encryptBlock :: (Integer, Integer) -> Integer -> Integer
encryptBlock (n, e) block = fastpow block e n

decryptBlock :: (Integer, Integer) -> Integer -> Integer
decryptBlock (n, d) block = fastpow block d n

main :: IO ()
main = do
  let message = "This is a small message."
  putStrLn ("Message was: " ++ message)

  let codePoints = map (toInteger . ord) message
  putStrLn ("\nCodePoints were: " ++ show codePoints)

  (pubKey, privKey) <- genKeys
  putStrLn ("\npubKey was: " ++ show pubKey)
  putStrLn ("\nprivKey was: " ++ show privKey)

  let encryptedPoints = map (encryptBlock pubKey) codePoints
  putStrLn ("\nEncrypted message: " ++ show encryptedPoints)

  let decryptedPoints = map (decryptBlock privKey) encryptedPoints
  putStrLn ("\nDecrypted points: " ++ show decryptedPoints)

  let decryptedText = map (chr . fromIntegral) decryptedPoints
  putStrLn ("\nDecrypted text: " ++ decryptedText)