import Data.Bits
import System.Random
import Data.Char

--Keys, in the form n, k (where k is i for the public key, and j for the private key)
data Key = Public Integer Integer | Private Integer Integer
	deriving (Eq, Ord, Show)

-- Extended Euclidian algorithm using the recursive method, returning (gcd, x, y)
eea :: (Integral a) => a -> a -> (a, a, a)
eea a b
	| b == 0	= (a, 1, 0)
	| otherwise	= (d, t, s - q * t)
	where
		(q, r) = a `divMod` b
		(d, s, t) = eea b r

-- Modular multiplicative inverse for a (mod m)
mminv :: (Integral a) => a -> a-> a
mminv a m
	| gcd /= 1	= error "Number doesn't have a multiplicative inverse for this modulus!"
	| otherwise	= x `mod` m
	where
		(gcd, x, _) = eea a m

-- Modular exponentiation by squaring (using Montgomery's ladder to prevent side-channel (i.e. implementation based)
-- attacks as discussed in http://www.sidechannelattacks.com/details/paper_details.aspx?fid=661
mexp :: Integer -> Integer -> Integer -> Integer
mexp x n m
	| n == 0	= 1
	| otherwise	= fst (foldl (mexp' m) (x, x ^ 2) [ testBit n (k - b - 2) | b <- [0 ..(k - 2)] ])
	where
		k = ceiling ( logBase 2 (fromIntegral (n + 1)) )

mexp' :: Integer -> (Integer, Integer) -> Bool -> (Integer, Integer)
mexp' m xs b
	| b == False	= ((x1 ^ 2) `mod` m, (x1 * x2) `mod` m)
	| otherwise		= ((x1 * x2) `mod` m, (x2 ^ 2) `mod` m)
	where
		x1 = fst xs
		x2 = snd xs

-- Generate public and private keys using the multiplicative inverse of i, mod phi
generateKeys :: Integer -> Integer -> Integer -> (Key, Key)
generateKeys p q i
	| gcd /= 1	= error "Public exponent i is not coprime with phi!"
	| otherwise	= (Public n i, Private n j)
	where
		n = p * q
		phi = (p - 1) * (q - 1)
		(gcd, _, _) = eea i phi
		j = mminv i phi

-- Code or decode an integer, given a public/private key
rsaCoder :: Key -> Integer -> Integer
rsaCoder (Public n k) x = mexp x k n
rsaCoder (Private n k) x = mexp x k n


-- Primality tester from http://www.haskell.org/haskellwiki/Testing_primality, but using my own mexp function instead of theirs
-- BEGIN --
-- (eq. to) find2km (2^k * n) = (k,n)
find2km :: Integral a => a -> (a,a)
find2km n = f 0 n
    where 
        f k m
            | r == 1 = (k,m)
            | otherwise = f (k+1) q
            where (q,r) = quotRem m 2  

-- n is the number to test; a is the (presumably randomly chosen) witness
millerRabinPrimality :: Integer -> Integer -> Bool
millerRabinPrimality n a
    | a <= 1 || a >= n-1 = 
        error $ "millerRabinPrimality: a out of range (" 
              ++ show a ++ " for "++ show n ++ ")" 
    | n < 2 = False
    | even n = False
    | b0 == 1 || b0 == n' = True
    | otherwise = iter (tail b)
    where
        n' = n-1
        (k,m) = find2km n'
        b0 = mexp a m n -- modified this line
        b = take (fromIntegral k) $ iterate (squareMod n) b0
        iter [] = False
        iter (x:xs)
            | x == 1 = False
            | x == n' = True
            | otherwise = iter xs

squareMod :: Integral a => a -> a -> a
squareMod a b = (b * b) `rem` a
-- END --

-- Use the Miller-Rabin method of primality testing, with a witness of 100 (i.e. a non-prime probability of 2^(-100),
-- according to http://snippets.dzone.com/posts/show/4200)
primeTest :: Integer -> Bool
primeTest x = millerRabinPrimality x 100

-- Generate an n-bit random prime number
getPrime :: Integer -> IO Integer
getPrime n = do
	r <- randomRIO (2 ^ n, (2 ^ (n + 1)) - 1)
	if (primeTest r)
		then
			return r
		else
			getPrime n

-- Encode a string byte-wise as an list of RSA-encrypted integers (this is not a good way of doing it, as frequency
-- analysis can easily be performed for frequent characters
encode:: String -> Key -> [Integer]
encode s k = [rsaCoder k (toInteger $ ord i) | i <- s]

-- Decode a list of RSA-encrypted integers byte-wise to a string
decode:: [Integer] -> Key -> String
decode s k = [chr $ fromInteger $ rsaCoder k i | i <- s]

main :: IO()
main = do
	p <- getPrime 256 
	q <- getPrime 256
	i <- getPrime 256
	putStr $ "p: " ++ (show p) ++ " " ++ show (primeTest p) ++ "\n"
	putStr $ "q: " ++ (show q) ++ " " ++ show (primeTest q) ++ "\n"
	putStr $ "i: " ++ (show i) ++ " " ++ show (primeTest i) ++ "\n"
	let keys = generateKeys p q i
	let pub = fst keys
	let priv = snd keys
	putStr $ show pub
	putStr "\n"
	putStr $ show priv
	putStr "\nType the text to encode:\n"
	plaintext <- getLine
	putStr "\n"
	let encrypted = encode plaintext pub
	putStr $ "Encrypted:\n" ++ (show encrypted) ++ "\n"
	let decrypted = decode encrypted priv
	putStr $ "Decrypted:\n" ++ (show decrypted) ++ "\n"