{-# LANGUAGE OverloadedStrings #-}

import Test.Framework (defaultMain, testGroup)
import Test.Framework.Providers.HUnit

import Test.HUnit

import Codec.Encryption.OpenPGP.Serialize ()
import Codec.Encryption.OpenPGP.ASCIIArmor (armor, decodeArmor)
import Codec.Encryption.OpenPGP.Compression (decompressPacket, compressPackets)
import Codec.Encryption.OpenPGP.Fingerprint (eightOctetKeyID, fingerprint)
import Codec.Encryption.OpenPGP.Types
import Data.Conduit.Cereal.Temp (conduitGet)
import Data.Conduit.OpenPGP.Compression (conduitCompress, conduitDecompress)
import Data.Conduit.OpenPGP.Keyring (conduitToTKs, sinkKeyringMap)
import Data.Conduit.OpenPGP.Verify (conduitVerify)

import Data.ByteString (ByteString)
import qualified Data.ByteString as B
import Data.Digest.CRC24 (crc24)
import qualified Data.Map as Map
import Data.Maybe (isJust)
import Data.Serialize (get, put)
import Data.Serialize.Get (runGet, Get)
import Data.Serialize.Put (runPut)
import Data.Word (Word32)

import qualified Data.Conduit as DC
import qualified Data.Conduit.Binary as CB
import qualified Data.Conduit.List as CL

testSerialization :: FilePath -> Assertion
testSerialization fp = do
    bs <- B.readFile $ "tests/data/" ++ fp
    let firstpass = runGet get bs
    case fmap unBlock firstpass of
        Left err -> assertFailure $ "First pass failed on " ++ fp
        Right [] -> assertFailure $ "First pass of " ++ fp ++ " decoded to nothing."
        Right packs -> do let roundtrip = runPut $ put (Block packs)
                          let secondpass = runGet (get :: Get (Block Packet)) roundtrip
                          if fmap unBlock secondpass == Right [] then
                              assertFailure $ "Second pass of " ++ fp ++ " decoded to nothing."
                          else
                              assertEqual ("for " ++ fp) firstpass secondpass

testCRC24 :: ByteString -> Word32 -> Assertion
testCRC24 bs crc = assertEqual "crc24" crc (crc24 bs)

testArmorDecode :: FilePath -> Armor Int -> Assertion
testArmorDecode fp target = do
    bs <- B.readFile $ "tests/data/" ++ fp
    case decodeArmor bs of
        Left err -> assertFailure $ "Decode failed on " ++ fp
        Right arm -> do assertEqual ("for " ++ fp) target arm

testArmorEncode :: Armor Int -> ByteString -> Assertion
testArmorEncode arm target = assertEqual ("literaldata") (armor arm) target

testCompression :: FilePath -> Assertion
testCompression fp = do
    bs <- B.readFile $ "tests/data/" ++ fp
    let firstpass = fmap (concatMap decompressPacket) . fmap unBlock . runGet get $ bs
    case firstpass of
        Left err -> assertFailure $ "First pass failed on " ++ fp
        Right [] -> assertFailure $ "First pass of " ++ fp ++ " decoded to nothing."
        Right packs -> do let roundtrip = runPut $ put . Block $ [compressPackets ZIP packs]
                          let secondpass = fmap (concatMap decompressPacket) . fmap unBlock . runGet get $ roundtrip
                          if secondpass == Right [] then
                              assertFailure $ "Second pass of " ++ fp ++ " decoded to nothing."
                          else
                              assertEqual ("for " ++ fp) firstpass secondpass

counter :: (DC.ResourceIO m) => DC.Sink a m Int
counter = CL.fold (const . (1+)) 0

testConduitOutputLength :: FilePath -> DC.Conduit B.ByteString IO b -> Int -> Assertion
testConduitOutputLength fp c target = do
    len <- DC.runResourceT $ CB.sourceFile ("tests/data/" ++ fp) DC.$= c DC.$$ counter
    assertEqual ("expected length" ++ show target) target len

testKeyIDandFingerprint :: FilePath -> String -> Assertion
testKeyIDandFingerprint fp kf = do
    bs <- B.readFile $ "tests/data/" ++ fp
    case runGet (get :: Get Packet) bs of
        Left err -> assertFailure $ "Decoding of " ++ fp ++ " broke."
        Right (PublicKey pkp) -> assertEqual ("for " ++ fp) kf (show (eightOctetKeyID pkp) ++ "/" ++ show (fingerprint pkp))

testKeyringLookup :: FilePath -> String -> Bool -> Assertion
testKeyringLookup fp eok expected = do
    kr <- DC.runResourceT $ CB.sourceFile ("tests/data/" ++ fp) DC.$= conduitGet get DC.$= conduitToTKs DC.$$ sinkKeyringMap
    let key = (Map.lookup (read eok) kr)
    assertEqual (eok ++ " in " ++ fp) expected (isJust key)

testVerifyMessage :: FilePath -> FilePath -> Assertion
testVerifyMessage keyring message = do
    kr <- DC.runResourceT $ CB.sourceFile ("tests/data/" ++ keyring) DC.$= conduitGet get DC.$= conduitToTKs DC.$$ sinkKeyringMap
    verification <- DC.runResourceT $ CB.sourceFile ("tests/data/" ++ message) DC.$= conduitGet get DC.$= conduitDecompress DC.$= conduitVerify kr DC.$$ CL.consume
    assertEqual (keyring ++ " for " ++ message) ([Right True]) verification

tests = [
  testGroup "Serialization group" [
     testCase "000001-006.public_key" (testSerialization "000001-006.public_key")
   , testCase "000002-013.user_id" (testSerialization "000002-013.user_id")
   , testCase "000003-002.sig" (testSerialization "000003-002.sig")
   , testCase "000004-012.ring_trust" (testSerialization "000004-012.ring_trust")
   , testCase "000005-002.sig" (testSerialization "000005-002.sig")
   , testCase "000006-012.ring_trust" (testSerialization "000006-012.ring_trust")
   , testCase "000007-002.sig" (testSerialization "000007-002.sig")
   , testCase "000008-012.ring_trust" (testSerialization "000008-012.ring_trust")
   , testCase "000009-002.sig" (testSerialization "000009-002.sig")
   , testCase "000010-012.ring_trust" (testSerialization "000010-012.ring_trust")
   , testCase "000011-002.sig" (testSerialization "000011-002.sig")
   , testCase "000012-012.ring_trust" (testSerialization "000012-012.ring_trust")
   , testCase "000013-014.public_subkey" (testSerialization "000013-014.public_subkey")
   , testCase "000014-002.sig" (testSerialization "000014-002.sig")
   , testCase "000015-012.ring_trust" (testSerialization "000015-012.ring_trust")
   , testCase "000016-006.public_key" (testSerialization "000016-006.public_key")
   , testCase "000017-002.sig" (testSerialization "000017-002.sig")
   , testCase "000018-012.ring_trust" (testSerialization "000018-012.ring_trust")
   , testCase "000019-013.user_id" (testSerialization "000019-013.user_id")
   , testCase "000020-002.sig" (testSerialization "000020-002.sig")
   , testCase "000021-012.ring_trust" (testSerialization "000021-012.ring_trust")
   , testCase "000022-002.sig" (testSerialization "000022-002.sig")
   , testCase "000023-012.ring_trust" (testSerialization "000023-012.ring_trust")
   , testCase "000024-014.public_subkey" (testSerialization "000024-014.public_subkey")
   , testCase "000025-002.sig" (testSerialization "000025-002.sig")
   , testCase "000026-012.ring_trust" (testSerialization "000026-012.ring_trust")
   , testCase "000027-006.public_key" (testSerialization "000027-006.public_key")
   , testCase "000028-002.sig" (testSerialization "000028-002.sig")
   , testCase "000029-012.ring_trust" (testSerialization "000029-012.ring_trust")
   , testCase "000030-013.user_id" (testSerialization "000030-013.user_id")
   , testCase "000031-002.sig" (testSerialization "000031-002.sig")
   , testCase "000032-012.ring_trust" (testSerialization "000032-012.ring_trust")
   , testCase "000033-002.sig" (testSerialization "000033-002.sig")
   , testCase "000034-012.ring_trust" (testSerialization "000034-012.ring_trust")
   , testCase "000035-006.public_key" (testSerialization "000035-006.public_key")
   , testCase "000036-013.user_id" (testSerialization "000036-013.user_id")
   , testCase "000037-002.sig" (testSerialization "000037-002.sig")
   , testCase "000038-012.ring_trust" (testSerialization "000038-012.ring_trust")
   , testCase "000039-002.sig" (testSerialization "000039-002.sig")
   , testCase "000040-012.ring_trust" (testSerialization "000040-012.ring_trust")
   , testCase "000041-017.attribute" (testSerialization "000041-017.attribute")
   , testCase "000042-002.sig" (testSerialization "000042-002.sig")
   , testCase "000043-012.ring_trust" (testSerialization "000043-012.ring_trust")
   , testCase "000044-014.public_subkey" (testSerialization "000044-014.public_subkey")
   , testCase "000045-002.sig" (testSerialization "000045-002.sig")
   , testCase "000046-012.ring_trust" (testSerialization "000046-012.ring_trust")
   , testCase "000047-005.secret_key" (testSerialization "000047-005.secret_key")
   , testCase "000048-013.user_id" (testSerialization "000048-013.user_id")
   , testCase "000049-002.sig" (testSerialization "000049-002.sig")
   , testCase "000050-012.ring_trust" (testSerialization "000050-012.ring_trust")
   , testCase "000051-007.secret_subkey" (testSerialization "000051-007.secret_subkey")
   , testCase "000052-002.sig" (testSerialization "000052-002.sig")
   , testCase "000053-012.ring_trust" (testSerialization "000053-012.ring_trust")
   , testCase "000054-005.secret_key" (testSerialization "000054-005.secret_key")
   , testCase "000055-002.sig" (testSerialization "000055-002.sig")
   , testCase "000056-012.ring_trust" (testSerialization "000056-012.ring_trust")
   , testCase "000057-013.user_id" (testSerialization "000057-013.user_id")
   , testCase "000058-002.sig" (testSerialization "000058-002.sig")
   , testCase "000059-012.ring_trust" (testSerialization "000059-012.ring_trust")
   , testCase "000060-007.secret_subkey" (testSerialization "000060-007.secret_subkey")
   , testCase "000061-002.sig" (testSerialization "000061-002.sig")
   , testCase "000062-012.ring_trust" (testSerialization "000062-012.ring_trust")
   , testCase "000063-005.secret_key" (testSerialization "000063-005.secret_key")
   , testCase "000064-002.sig" (testSerialization "000064-002.sig")
   , testCase "000065-012.ring_trust" (testSerialization "000065-012.ring_trust")
   , testCase "000066-013.user_id" (testSerialization "000066-013.user_id")
   , testCase "000067-002.sig" (testSerialization "000067-002.sig")
   , testCase "000068-012.ring_trust" (testSerialization "000068-012.ring_trust")
   , testCase "000069-005.secret_key" (testSerialization "000069-005.secret_key")
   , testCase "000070-013.user_id" (testSerialization "000070-013.user_id")
   , testCase "000071-002.sig" (testSerialization "000071-002.sig")
   , testCase "000072-012.ring_trust" (testSerialization "000072-012.ring_trust")
   , testCase "000073-017.attribute" (testSerialization "000073-017.attribute")
   , testCase "000074-002.sig" (testSerialization "000074-002.sig")
   , testCase "000075-012.ring_trust" (testSerialization "000075-012.ring_trust")
   , testCase "000076-007.secret_subkey" (testSerialization "000076-007.secret_subkey")
   , testCase "000077-002.sig" (testSerialization "000077-002.sig")
   , testCase "000078-012.ring_trust" (testSerialization "000078-012.ring_trust")
   , testCase "pubring.gpg" (testSerialization "pubring.gpg")
   , testCase "secring.gpg" (testSerialization "secring.gpg")
   , testCase "compressedsig.gpg" (testSerialization "compressedsig.gpg")
   , testCase "compressedsig-zlib.gpg" (testSerialization "compressedsig-zlib.gpg")
   , testCase "compressedsig-bzip2.gpg" (testSerialization "compressedsig-bzip2.gpg")
   , testCase "onepass_sig" (testSerialization "onepass_sig")
   , testCase "uncompressed-ops-dsa.gpg" (testSerialization "uncompressed-ops-dsa.gpg")
   , testCase "uncompressed-ops-rsa.gpg" (testSerialization "uncompressed-ops-rsa.gpg")
   ],
  testGroup "CRC24 group" [
     testCase "CRC24: A" (testCRC24 "A" 16680698)
   , testCase "CRC24: Haskell" (testCRC24 "Haskell" 15612750)
   , testCase "CRC24: hOpenPGP and friends" (testCRC24 "hOpenPGP and friends" 11940960)
   ],
  testGroup "ASCII armor group" [
     testCase "Decode sample armor" (testArmorDecode "msg1.asc" (Armor ArmorMessage [("Version","OpenPrivacy 0.99")] [CompressedData ZIP ";m\150\196\DC1\239\236\239\ETB\236\239\227\202\NUL\EOT\206\137y\234%\n\137y\149\249y\169\n\217\169\169\ENQ\n\137\n\197\169\201E\169@\193\162\252\210\188\DC4\133\140\212\162T{.\NUL"]))
   , testCase "Encode sample armor" (testArmorEncode (Armor ArmorMessage [("Comment", "Test")] [LiteralData UTF8Data "notlob.txt" 12345 "These are the days of literal data.\n"]) "-----BEGIN PGP MESSAGE-----\nComment: Test\n\nyzR1Cm5vdGxvYi50eHQAADA5VGhlc2UgYXJlIHRoZSBkYXlzIG9mIGxpdGVyYWwgZGF0YS4K\n=AAKY\n-----END PGP MESSAGE-----\n")
   ],
  testGroup "Compression group" [
     testCase "compressedsig.gpg" (testCompression "compressedsig.gpg")
   , testCase "compressedsig-zlib.gpg" (testCompression "compressedsig-zlib.gpg")
   , testCase "compressedsig-bzip2.gpg" (testCompression "compressedsig-bzip2.gpg")
   ],
  testGroup "Conduit group" [
     testCase "compressedsig straight" (testConduitOutputLength "compressedsig.gpg" (conduitGet (get :: Get Packet)) 1)
   , testCase "compressedsig uncompressed" (testConduitOutputLength "compressedsig.gpg" (conduitGet (get :: Get Packet) DC.=$= conduitDecompress) 3)
   , testCase "pubring" (testConduitOutputLength "pubring.gpg" (conduitGet (get :: Get Packet)) 46)
   ],
  testGroup "Fingerprint group" [
     testCase "000001-006.public_key" (testKeyIDandFingerprint "000001-006.public_key" "D4D54EA16F87040E/421F 28FE AAD2 22F8 56C8  FFD5 D4D5 4EA1 6F87 040E")
   , testCase "000016-006.public_key" (testKeyIDandFingerprint "000016-006.public_key" "5E9F1523413262DC/AF95 E4D7 BAC5 21EE 9740  BED7 5E9F 1523 4132 62DC")
   , testCase "000027-006.public_key" (testKeyIDandFingerprint "000027-006.public_key" "7732CF988A63EA86/1EB2 0B2F 5A5C C3BE AFD6  E5CB 7732 CF98 8A63 EA86")
   , testCase "000035-006.public_key" (testKeyIDandFingerprint "000035-006.public_key" "DEDC3ECF689AF56D/CB79 3345 9F59 C70D F1C3  FBEE DEDC 3ECF 689A F56D")
   ],
  testGroup "Keyring group" [
     testCase "pubring 7732CF988A63EA86" (testKeyringLookup "pubring.gpg" "7732CF988A63EA86" True)
   , testCase "pubring 123456789ABCDEF0" (testKeyringLookup "pubring.gpg" "123456789ABCDEF0" False)
   , testCase "secring 7732CF988A63EA86" (testKeyringLookup "secring.gpg" "7732CF988A63EA86" True)
   , testCase "secring 123456789ABCDEF0" (testKeyringLookup "secring.gpg" "123456789ABCDEF0" False)
   -- FIXME: should count keys in rings
   ],
  testGroup "Message verification group" [
     testCase "uncompressed-ops-dsa" (testVerifyMessage "pubring.gpg" "uncompressed-ops-dsa.gpg")
   , testCase "uncompressed-ops-dsa-sha384" (testVerifyMessage "pubring.gpg" "uncompressed-ops-dsa-sha384.txt.gpg")
   -- , testCase "uncompressed-ops-rsa" (testVerifyMessage "pubring.gpg" "uncompressed-ops-rsa.gpg")
   -- , testCase "compressedsig" (testVerifyMessage "pubring.gpg" "compressedsig.gpg")
   -- , testCase "compressedsig-zlib" (testVerifyMessage "pubring.gpg" "compressedsig-zlib.gpg")
   -- , testCase "compressedsig-bzip2" (testVerifyMessage "pubring.gpg" "compressedsig-bzip2.gpg")
   ]
 ]

main = defaultMain tests