module Main where

import System.IO
import System.Environment
import Control.Exception
import qualified Data.Maybe as Maybe
import qualified CyclicGroup as CG
import qualified ElGamal as EG
import qualified Encoder as Enc
import qualified NumberTheory as NT
import qualified Random as Rnd

-- Each command line action dispatches a function call generating I/O
dispatch :: [(String, [String] -> IO())]
dispatch = [ ("init", Main.init),
             ("keys", Main.keys),
             ("encrypt", Main.encrypt),
             ("decrypt", Main.decrypt) ]

print_usage :: IOError -> IO()
print_usage err = do
    putStrLn "init n       \t- Create an n-bit ElGamal cryptosystem"
    putStrLn "keys s       \t- Create keys within cryptosystem s"
    putStrLn "encrypt s k m\t- Encypt message m with public key k in system s"
    putStrLn "decrypt s p c\t- Decrypt cipher c with private key p in system s"

-- Create an ElGamal crypto system, i.e. a group
-- of prime order p and a generator g
init :: [String] -> IO()
init [num_bits] = do 
    start <- Rnd.random_nbit_int (read num_bits)
    let (p, q) = NT.next_pq_prime start
    let g = NT.next_primitive_root p q (div p 2)
    putStrLn ((show p) ++ " " ++ (show g))

-- Create private/public key pair given an ElGamal
-- crypto system (maybe generated by Main.init).
keys :: [String] -> IO()
keys [order, generator] = do
     priv_key <- Rnd.random_int (read order)
     let pub_key = NT.fast_pow (read generator) priv_key (read order)
     putStrLn ((show priv_key) ++ " " ++ (show pub_key))

-- Encrypt a message (an Integer) less than the order of the
-- cryptosystem.
-- FIXME nicer parameters to EG.encrypt
-- FIXME reusable code that makes random integer
encrypt :: [String] -> IO()
encrypt [order, generator, pub_key, message] = do
    let group = CG.Group (read order)
    let g = CG.Element (read generator)
    let pub = CG.Element (read pub_key)
    let msg = CG.Element (read message)
    -- Now we need an ephemeral key tmp
    tmp <- Rnd.random_int (read order)
    -- Finally we can generate the ElGamal cipher
    let (CG.Element c1, CG.Element c2) = EG.encrypt group msg g tmp pub
    putStrLn ((show c1) ++ " " ++ (show c2))


-- Encrypt a message (an Integer) less than the order of the
-- cryptosystem.
-- FIXME nicer parameters to EG.encrypt
-- FIXME reusable code that makes random integer
decrypt :: [String] -> IO()
decrypt [order, generator, priv_key, c1_str, c2_str] = do
    let group = CG.Group (read order)
    let g = CG.Element (read generator)
    let priv = read priv_key
    let (c1, c2) = (CG.Element (read c1_str), CG.Element (read c2_str))
    let msg = EG.decrypt group (c1, c2) priv
    putStrLn (show msg)

start = do
         (command:args) <- getArgs
         let (Just action) = lookup command dispatch
         action args

main = start `catch` print_usage

