{-# LANGUAGE OverloadedStrings #-}

-- suite.hs: hOpenPGP test suite
-- Copyright © 2012-2013  Clint Adams
-- This software is released under the terms of the Expat license.
-- (See the LICENSE file).

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

import Test.HUnit (Assertion, assertFailure, assertEqual)

import Codec.Encryption.OpenPGP.Serialize ()
import Codec.Encryption.OpenPGP.Compression (decompressPkt, compressPkts)
import Codec.Encryption.OpenPGP.Fingerprint (eightOctetKeyID, fingerprint)
import Codec.Encryption.OpenPGP.Types
import Data.Conduit.Cereal (conduitGet)
import Data.Conduit.OpenPGP.Compression (conduitCompress, conduitDecompress)
import Data.Conduit.OpenPGP.Decrypt (conduitDecrypt)
import Data.Conduit.OpenPGP.Keyring (conduitToTKs, conduitToTKsDropping, sinkKeyringMap)
import Data.Conduit.OpenPGP.Verify (conduitVerify)
import Codec.Encryption.OpenPGP.Serialize ()

import qualified Data.ByteString as B
import qualified Data.ByteString.Lazy as BL
import Data.IxSet ((@=), getOne)
import Data.Maybe (isJust)
import Data.Serialize (get, put)
import Data.Serialize.Get (runGet, Get)
import Data.Serialize.Put (runPut)

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

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

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

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

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

testKeyIDandFingerprint :: FilePath -> String -> Assertion
testKeyIDandFingerprint fpr kf = do
    bs <- B.readFile $ "tests/data/" ++ fpr
    case runGet (get :: Get Pkt) bs of
        Left _ -> assertFailure $ "Decoding of " ++ fpr ++ " broke."
        Right (PublicKeyPkt pkp) -> assertEqual ("for " ++ fpr) kf (show (eightOctetKeyID pkp) ++ "/" ++ show (fingerprint pkp))
        _ -> assertFailure "Expected public key, got something else."

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

testVerifyMessage :: FilePath -> FilePath -> [TwentyOctetFingerprint] -> Assertion
testVerifyMessage keyring message issuers = 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 Nothing DC.$$ CL.consume
    let verification' = map (fmap (fingerprint . _verificationSigner)) verification
    assertEqual (keyring ++ " for " ++ message) (map Right issuers) verification'

-- This needs a lot of work
testSymmetricEncryption :: FilePath -> FilePath -> B.ByteString -> Assertion
testSymmetricEncryption encfile passfile cleartext = do
  passphrase <- BL.readFile $ "tests/data/" ++ passfile
  -- get parse tree
  pt <- DC.runResourceT $ CB.sourceFile ("tests/data/" ++ encfile) DC.$= conduitGet get DC.$$ CL.consume
  -- assert parse tree has exactly two packets: skesk, encdata
  assertEqual "wrong number of packets" 2 (length pt)
  let skesk = fromPkt.head $ pt
      d = fromPkt.last $ pt
  -- FIXME: these assertions don't currently do anything properly,
  -- because haskell notices the _-prefixed accessor invocations below
  -- and the type system chokes before we hit them:
  assertEqual "first packet should be SKESK" SKESKType (packetType skesk)
  assertEqual "second packet should be encrypted data" SymEncIntegrityProtectedDataType (packetType d)
  decrypted <- DC.runResourceT $ CL.sourceList pt DC.$= conduitDecrypt (fakeCallback passphrase) DC.$$ CL.consume
  let payload = _literalDataPayload . fromPkt . head $ decrypted
  assertEqual ("cleartext for " ++ encfile) cleartext payload
      where
          fakeCallback :: BL.ByteString -> String -> IO BL.ByteString
          fakeCallback = const . return

tests :: [Test]
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")
   , testCase "simple.seckey" (testSerialization "simple.seckey")
   ],
  testGroup "KeyID/fingerprint group" [
     testCase "v3 key" (testKeyIDandFingerprint "v3.key" "C7261095/CBD9 F412 6807 E405 CC2D  2712 1DF5 E86E  ")
   , testCase "v4 key" (testKeyIDandFingerprint "000001-006.public_key" "D4D54EA16F87040E/421F 28FE AAD2 22F8 56C8  FFD5 D4D5 4EA1 6F87 040E")
   ],
  testGroup "Keyring group" [
     testCase "pubring 7732CF988A63EA86" (testKeyringLookup "pubring.gpg" "7732CF988A63EA86" True)
   , testCase "pubring 123456789ABCDEF0" (testKeyringLookup "pubring.gpg" "123456789ABCDEF0" False)
   , testCase "pubsub AD992E9C24399832" (testKeyringLookup "pubring.gpg" "AD992E9C24399832" True)
   , testCase "secring 7732CF988A63EA86" (testKeyringLookup "secring.gpg" "7732CF988A63EA86" True)
   , testCase "secring 123456789ABCDEF0" (testKeyringLookup "secring.gpg" "123456789ABCDEF0" False)
   , testCase "secsub AD992E9C24399832" (testKeyringLookup "secring.gpg" "AD992E9C24399832" True)
   -- FIXME: should count keys in rings
   ],
  testGroup "Message verification group" [
     testCase "uncompressed-ops-dsa" (testVerifyMessage "pubring.gpg" "uncompressed-ops-dsa.gpg" ([fp "1EB2 0B2F 5A5C C3BE AFD6  E5CB 7732 CF98 8A63 EA86"]))
   , testCase "uncompressed-ops-dsa-sha384" (testVerifyMessage "pubring.gpg" "uncompressed-ops-dsa-sha384.txt.gpg" ([fp "1EB2 0B2F 5A5C C3BE AFD6  E5CB 7732 CF98 8A63 EA86"]))
   , testCase "uncompressed-ops-rsa" (testVerifyMessage "pubring.gpg" "uncompressed-ops-rsa.gpg" ([fp "CB79 3345 9F59 C70D F1C3  FBEE DEDC 3ECF 689A F56D"]))
   , testCase "compressedsig" (testVerifyMessage "pubring.gpg" "compressedsig.gpg" ([fp "421F 28FE AAD2 22F8 56C8  FFD5 D4D5 4EA1 6F87 040E"]))
   , testCase "compressedsig-zlib" (testVerifyMessage "pubring.gpg" "compressedsig-zlib.gpg" ([fp "421F 28FE AAD2 22F8 56C8  FFD5 D4D5 4EA1 6F87 040E"]))
   , testCase "compressedsig-bzip2" (testVerifyMessage "pubring.gpg" "compressedsig-bzip2.gpg" ([fp "421F 28FE AAD2 22F8 56C8  FFD5 D4D5 4EA1 6F87 040E"]))
   ],
  testGroup "Certificate verification group" [
     testCase "userid" (testVerifyMessage "pubring.gpg" "minimized.gpg" ([fp "421F 28FE AAD2 22F8 56C8  FFD5 D4D5 4EA1 6F87 040E"]))
   , testCase "subkey" (testVerifyMessage "pubring.gpg" "subkey.gpg" ([fp "421F 28FE AAD2 22F8 56C8  FFD5 D4D5 4EA1 6F87 040E"]))
   , testCase "primary key binding" (testVerifyMessage "signing-subkey.gpg" "primary-binding.gpg" ([fp "ED1B D216 F70E 5D5F 4444  48F9 B830 F2C4 83A9 9AE5"]))
   , testCase "attribute" (testVerifyMessage "pubring.gpg" "uat.gpg" ([fp "421F 28FE AAD2 22F8 56C8  FFD5 D4D5 4EA1 6F87 040E"]))
   , testCase "primary key revocation" (testVerifyMessage "pubring.gpg" "prikey-rev.gpg" ([fp "421F 28FE AAD2 22F8 56C8  FFD5 D4D5 4EA1 6F87 040E"]))
   , testCase "subkey revocation" (testVerifyMessage "pubring.gpg" "subkey-rev.gpg" ([fp "421F 28FE AAD2 22F8 56C8  FFD5 D4D5 4EA1 6F87 040E"]))
   , testCase "6F87040E" (testVerifyMessage "pubring.gpg" "6F87040E.pubkey" ([fp "421F 28FE AAD2 22F8 56C8  FFD5 D4D5 4EA1 6F87 040E", fp "CB79 3345 9F59 C70D F1C3  FBEE DEDC 3ECF 689A F56D", fp "AF95 E4D7 BAC5 21EE 9740  BED7 5E9F 1523 4132 62DC"]))
   , testCase "6F87040E-cr" (testVerifyMessage "pubring.gpg" "6F87040E-cr.pubkey" ([fp "AF95 E4D7 BAC5 21EE 9740  BED7 5E9F 1523 4132 62DC", fp "AF95 E4D7 BAC5 21EE 9740  BED7 5E9F 1523 4132 62DC", fp "421F 28FE AAD2 22F8 56C8  FFD5 D4D5 4EA1 6F87 040E", fp "CB79 3345 9F59 C70D F1C3  FBEE DEDC 3ECF 689A F56D", fp "AF95 E4D7 BAC5 21EE 9740  BED7 5E9F 1523 4132 62DC"]))
   , testCase "simple RSA secret key" (testVerifyMessage "pubring.gpg" "simple.seckey" ([fp "421F 28FE AAD2 22F8 56C8  FFD5 D4D5 4EA1 6F87 040E"]))
   ],
  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 length group" [
     testCase "conduitCompress (ZIP)" (testConduitOutputLength "pubring.gpg" (cgp DC.=$= conduitCompress ZIP) 1)
   , testCase "conduitCompress (Zlib)" (testConduitOutputLength "pubring.gpg" (cgp DC.=$= conduitCompress ZLIB) 1)
   , testCase "conduitCompress (BZip2)" (testConduitOutputLength "pubring.gpg" (cgp DC.=$= conduitCompress BZip2) 1)
   , testCase "conduitToTKs" (testConduitOutputLength "pubring.gpg" (cgp DC.=$= conduitToTKs) 4)
   , testCase "conduitToTKsDropping" (testConduitOutputLength "pubring.gpg" (cgp DC.=$= conduitToTKsDropping) 4)
   ],
  testGroup "Encrypted data" [
     testCase "Symmetric Encryption simple S2K SHA1 3DES, no MDC" (testSymmetricEncryption "encryption-sym-3des-s2k0.gpg" "symmetric-password.txt" "test\n")
   , testCase "Symmetric Encryption iterated-salted S2K SHA1 3DES, no MDC" (testSymmetricEncryption "encryption-sym-3des.gpg" "symmetric-password.txt" "test\n")
   , testCase "Symmetric Encryption simple S2K SHA1 3DES" (testSymmetricEncryption "encryption-sym-3des-mdc-s2k0.gpg" "symmetric-password.txt" "test\n")
   , testCase "Symmetric Encryption iterated-salted S2K SHA1 3DES" (testSymmetricEncryption "encryption-sym-3des-mdc.gpg" "symmetric-password.txt" "test\n")
   , testCase "Symmetric Encryption simple S2K SHA1 CAST5, no MDC" (testSymmetricEncryption "encryption-sym-cast5-s2k0.gpg" "symmetric-password.txt" "test\n")
   , testCase "Symmetric Encryption iterated-salted S2K SHA1 CAST5, no MDC" (testSymmetricEncryption "encryption-sym-cast5.gpg" "symmetric-password.txt" "test\n")
   , testCase "Symmetric Encryption simple S2K SHA1 CAST5" (testSymmetricEncryption "encryption-sym-cast5-mdc-s2k0.gpg" "symmetric-password.txt" "test\n")
   , testCase "Symmetric Encryption iterated-salted S2K SHA1 CAST5" (testSymmetricEncryption "encryption-sym-cast5-mdc.gpg" "symmetric-password.txt" "test\n")
   , testCase "Symmetric Encryption simple S2K SHA1 Blowfish, no MDC" (testSymmetricEncryption "encryption-sym-blowfish-s2k0.gpg" "symmetric-password.txt" "test\n")
   , testCase "Symmetric Encryption iterated-salted S2K SHA1 Blowfish, no MDC" (testSymmetricEncryption "encryption-sym-blowfish.gpg" "symmetric-password.txt" "test\n")
   , testCase "Symmetric Encryption simple S2K SHA1 Blowfish" (testSymmetricEncryption "encryption-sym-blowfish-mdc-s2k0.gpg" "symmetric-password.txt" "test\n")
   , testCase "Symmetric Encryption iterated-salted S2K SHA1 Blowfish" (testSymmetricEncryption "encryption-sym-blowfish-mdc.gpg" "symmetric-password.txt" "test\n")
   , testCase "Symmetric Encryption simple S2K SHA1 AES128" (testSymmetricEncryption "encryption-sym-aes128-s2k0.gpg" "symmetric-password.txt" "test\n")
   , testCase "Symmetric Encryption iterated-salted S2K SHA1 AES128" (testSymmetricEncryption "encryption-sym-aes128.gpg" "symmetric-password.txt" "test\n")
   , testCase "Symmetric Encryption simple S2K SHA1 AES192" (testSymmetricEncryption "encryption-sym-aes192-s2k0.gpg" "symmetric-password.txt" "test\n")
   , testCase "Symmetric Encryption iterated-salted S2K SHA1 AES192" (testSymmetricEncryption "encryption-sym-aes192.gpg" "symmetric-password.txt" "test\n")
   , testCase "Symmetric Encryption simple S2K SHA1 AES256" (testSymmetricEncryption "encryption-sym-aes256-s2k0.gpg" "symmetric-password.txt" "test\n")
   , testCase "Symmetric Encryption iterated-salted S2K SHA1 AES256" (testSymmetricEncryption "encryption-sym-aes256.gpg" "symmetric-password.txt" "test\n")
   , testCase "Symmetric Encryption simple S2K SHA1 Twofish" (testSymmetricEncryption "encryption-sym-twofish-s2k0.gpg" "symmetric-password.txt" "test\n")
   , testCase "Symmetric Encryption iterated-salted S2K SHA1 Twofish" (testSymmetricEncryption "encryption-sym-twofish.gpg" "symmetric-password.txt" "test\n")
   ]
 ]

cgp :: DC.Conduit B.ByteString (DC.ResourceT IO) Pkt
cgp = conduitGet (get :: Get Pkt)

fp :: String -> TwentyOctetFingerprint
fp = read

main :: IO ()
main = defaultMain tests