module RSA where

import System.IO.Unsafe
import System.Random

erweiterterEuklid :: Integral a => a -> a -> (a,a,a)
erweiterterEuklid 0 y = (0,1,y)
erweiterterEuklid x y = (a,b,g) -- Invariant a*x + b*y = g
  where (a',b',g) = erweiterterEuklid x' y'
        x' = y `mod` x
	y' = x
	a  = b' - b*(y `div` x)
	b  = a'

modInvers :: Integral a => a -> a -> [a]
modInvers x m = case erweiterterEuklid x m of
                  (a, _, 1) -> [modPositiv a m]
		  (_, _, _) -> []

modPositiv :: Integral a => a -> a -> a
modPositiv x m | x < 0     = m + x
               | otherwise = x

factors :: Integer -> [Integer]
factors n = [x | x <- takeWhile (\p -> p*p <= n) prims,
                 n `mod` x == 0]

isPrim :: Integer -> Bool
isPrim = null . factors

prims :: [Integer]
prims = 2 : filter isPrim [3, 5 ..]

miniPrimProduct :: Int -> Integer
miniPrimProduct n = miniPrimProducts !! n

miniPrimProducts :: [Integer]
miniPrimProducts  = tail $ scanl (*) 1 prims

isMiniPrim :: Int -> Integer -> Bool
isMiniPrim n x = 1 == x `gcd` miniPrimProduct n

randomNumberIO :: Int -> IO Integer
randomNumberIO bits = randomRIO (2^(bits-1), 2^(bits)-1)

randomNumber :: Int -> Integer
randomNumber bits = unsafePerformIO (randomNumberIO bits)

getRandoms :: Random a => (a,a) -> [a]
getRandoms bnd = unsafePerformIO (newStdGen >>= return . randomRs bnd)

getWitnesses :: (Num a, Random a) => Int -> a -> [a]
getWitnesses n x = take n $ getRandoms (2, x-1)

primWith :: (Random a, Num a) => (a -> Bool) -> Int -> a -> Bool
primWith f n x = all f $ getWitnesses n x

fermatTest :: Integer -> Integer -> Bool
fermatTest x = validHead . rabinMillerSeq x
  where validHead (1:_) = True
        validHead _     = False

isFermatPrim :: Int -> Integer -> Bool
isFermatPrim n x = primWith (fermatTest x) n x

lehmanTest :: Integer -> Integer -> Bool
lehmanTest x = validHead . rabinMillerSeq x
  where validHead (_:y:_) = y `elem` [1, x - 1]
        validHead _       = False

isLehmanPrim :: Int -> Integer -> Bool
isLehmanPrim n x = primWith (lehmanTest x) n x

rabinMillerSeq :: Integer -> Integer -> [Integer]
rabinMillerSeq x r = map (\e -> expMod r e x) .
                     takeWhile (0 /=) .
		     iterate (\e -> if even e then e `div` 2 else 0) $
		     (x-1)

rabinMillerTest :: Integer -> Integer -> Bool
rabinMillerTest x = validHead . dropWhile (1 ==) . rabinMillerSeq x
  where validHead (y:_) = y == x - 1
        validHead []    = True

isRabinMillerPrim :: Int -> Integer -> Bool
isRabinMillerPrim n x = primWith (rabinMillerTest x) n x

filters :: [a -> Bool] -> ([a] -> [a])
filters = foldr (.) id . map filter

odds :: Integral a => a -> [a]
odds x = [ start, start + 2 .. ]
     where start = if odd x then x else x + 1

nextPrims :: Integer -> [Integer]
nextPrims = filters [isRabinMillerPrim 6, isMiniPrim 10] . odds

nextPrim :: Integer -> Integer
nextPrim = head . nextPrims

nextPrimB :: Int -> Integer
nextPrimB bits = nextPrim (1 + 2 * randomNumber (bits-1))

class Ideal i where
  expMod :: Integer -> Integer -> i -> Integer

instance Ideal Integer where
  expMod x e m = em (mm x 1) e
       where mm x y = (x*y) `mod` m
             em x 0 = 1
             em x e | even e    =        em (mm x x) (e `div` 2)
	     	    | otherwise = mm x $ em (mm x x) (e `div` 2)

class Ideal k => Key k where
  defaultExponent :: k -> Integer
  crypt :: k -> Integer -> Integer
  crypt k x = expMod x (defaultExponent k) k

data PubKey = PubKey { m, e :: Integer }
  deriving (Show, Eq)

instance Ideal PubKey where
  expMod x e k = expMod x e (m k)

instance Key PubKey where
  defaultExponent = e

data PrivKey = PrivKey { d, p, q, u :: Integer }
  deriving (Show, Eq)

instance Ideal PrivKey where
  -- Chinesischer Restsatz
  expMod x e k = xp + pk * (((xq - xp) * uk) `mod` qk)
      where xp = expMod x (e `mod` (pk - 1)) pk
            xq = expMod x (e `mod` (qk - 1)) qk
	    pk = p k
	    uk = u k
	    qk = q k

instance Key PrivKey where
  defaultExponent = d

data KeyPair = KeyPair { public :: PubKey, private :: PrivKey }
  deriving (Show, Eq)

instance Ideal KeyPair where
  expMod x e k = expMod x e (private k)

genRSA :: Integer -> Int -> KeyPair
genRSA e bits = head [ KeyPair (PubKey  {e=e, m=p*q})
                               (PrivKey {d=d, p=p, q=q, u=u})
                     | let p1    = nextPrimB (bits `div` 2),
		       let start = 3 * 2^(bits - 2) + randomNumber (bits - 2),
                       p2       <- nextPrims (start `div` p1),
		       let (p,q) = (min p1 p2, max p1 p2),
		       u        <- modInvers p q,
		       let phi   = (p-1)*(q-1),
		       d   	<- modInvers e phi
	             ]

stdRSA :: Int -> KeyPair
stdRSA = genRSA (2^16+1)