import Bindings.Nettle.Cipher.CAST128
import qualified Data.ByteString as B
import Control.Monad (foldM, replicateM)
import Data.Word (Word8)
import Foreign.Marshal.Alloc
import Foreign.Marshal.Array (withArray, peekArray)

import Test.Hspec
import Test.QuickCheck
import Test.QuickCheck.IO ()

plaintext :: [Word8]
plaintext = [0x01,0x23,0x45,0x67,0x89,0xAB,0xCD,0xEF]

key128 :: [Word8]
key128 = [0x01,0x23,0x45,0x67,0x12,0x34,0x56,0x78,0x23,0x45,0x67,0x89,0x34,0x56,0x78,0x9A]

key80 :: [Word8]
key80 = [0x01,0x23,0x45,0x67,0x12,0x34,0x56,0x78,0x23,0x45]

key40 :: [Word8]
key40 = [0x01,0x23,0x45,0x67,0x12]

ciphertext128 :: [Word8]
ciphertext128 = [0x23,0x8B,0x4F,0xE5,0x84,0x7E,0x44,0xB2]

ciphertext80 :: [Word8]
ciphertext80 = [0xEB,0x6A,0x71,0x1A,0x2C,0x02,0x27,0x1B]

ciphertext40 :: [Word8]
ciphertext40 = [0x7A,0xC8,0x16,0xD1,0x6E,0x9B,0x30,0x2E]

cast5_decrypt :: B.ByteString -> B.ByteString -> IO B.ByteString
cast5_decrypt key payload = alloca $ \ctx -> allocaBytes (B.length payload) $ \buf -> withArray (map fromIntegral . B.unpack $ key) $ \k -> withArray (map fromIntegral . B.unpack $ payload) $ \src -> do
    _ <- c'nettle_cast128_set_key ctx (fromIntegral . B.length $ key) k
    _ <- c'nettle_cast128_decrypt ctx (fromIntegral . B.length $ payload) buf src
    res <- peekArray (B.length payload) buf
    return $ B.pack $ map fromIntegral res

cast5_encrypt :: B.ByteString -> B.ByteString -> IO B.ByteString
cast5_encrypt key payload = alloca $ \ctx -> allocaBytes (B.length payload) $ \buf -> withArray (map fromIntegral . B.unpack $ key) $ \k -> withArray (map fromIntegral . B.unpack $ payload) $ \src -> do
    _ <- c'nettle_cast128_set_key ctx (fromIntegral . B.length $ key) k
    _ <- c'nettle_cast128_encrypt ctx (fromIntegral . B.length $ payload) buf src
    res <- peekArray (B.length payload) buf
    return $ B.pack $ map fromIntegral res

fmtInitial :: B.ByteString
fmtInitial = B.pack [01,0x23,0x45,0x67,0x12,0x34,0x56,0x78,0x23,0x45,0x67,0x89,0x34,0x56,0x78,0x9A]

fmtOutput :: (B.ByteString, B.ByteString)
fmtOutput = (B.pack [0xEE,0xA9,0xD0,0xA2,0x49,0xFD,0x3B,0xA6,0xB3,0x43,0x6F,0xB8,0x9D,0x6D,0xCA,0x92]
            ,B.pack [0xB2,0xC9,0x5E,0xB0,0x0C,0x31,0xAD,0x71,0x80,0xAC,0x05,0xB8,0xE8,0x3D,0x69,0x6E])

fmt :: IO (B.ByteString, B.ByteString)
fmt = stupidM 1000000 fmt' (fmtInitial, fmtInitial)
    where
        fmt' (a,b) = do
            al <- cast5_encrypt b (B.take 8 a)
            ar <- cast5_encrypt b (B.drop 8 a)
            let a' = B.append al ar
            bl <- cast5_encrypt a' (B.take 8 b)
            br <- cast5_encrypt a' (B.drop 8 b)
            let b' = B.append bl br
            return (a', b')
        stupidM n f iv = foldM (const . f) iv [1..n]

prop_cast5_roundtrip :: Property
prop_cast5_roundtrip = forAll cast5_keys $ \key -> forAll cast5_plaintexts $ \pt -> do
  c <- cast5_encrypt key pt
  cast5_decrypt key c `shouldReturn` pt

cast5_keys :: Gen B.ByteString
cast5_keys = do
    l <- fmap ((+5) . (`mod` 12) . abs) arbitrary
    os <- replicateM l arbitrary
    return $ B.pack os

cast5_plaintexts :: Gen B.ByteString
cast5_plaintexts = do
    l <- fmap ((*8) . (`mod` 100) . abs) arbitrary
    os <- replicateM l arbitrary
    return $ B.pack os


main :: IO ()
main = hspec spec

spec :: Spec
spec = do
    describe "CAST128 decryption" $ do
       it "matches RFC example with 128-bit key" $ cast5_decrypt (B.pack key128) (B.pack ciphertext128) `shouldReturn` B.pack plaintext
       it "matches RFC example with 80-bit key" $ cast5_decrypt (B.pack key80) (B.pack ciphertext80) `shouldReturn` B.pack plaintext
       it "matches RFC example with 40-bit key" $ cast5_decrypt (B.pack key40) (B.pack ciphertext40) `shouldReturn` B.pack plaintext
       it "is the inverse of CAST128 encryption" $ property $ prop_cast5_roundtrip
    describe "CAST128 encryption" $ do
       it "matches RFC example with 128-bit key" $ cast5_encrypt (B.pack key128) (B.pack plaintext) `shouldReturn` B.pack ciphertext128
       it "matches RFC example with 80-bit key" $ cast5_encrypt (B.pack key80) (B.pack plaintext) `shouldReturn` B.pack ciphertext80
       it "matches RFC example with 40-bit key" $ cast5_encrypt (B.pack key40) (B.pack plaintext) `shouldReturn` B.pack ciphertext40
       it "matches RFC FMT example" $ fmt `shouldReturn` fmtOutput