{-# LANGUAGE OverloadedStrings #-}

module Main (main) where

import           Data.ByteString        (ByteString)

import qualified Data.ByteString        as B

import qualified Data.ByteString.Base16 as B16

import qualified Data.ByteString.Lazy   as BL

import           Data.Word

-- reference implementation

import qualified Data.Digest.Pure.SHA   as REF

-- implementation under test

import qualified Crypto.Hash.SHA256     as IUT

import           Test.Tasty

import           Test.Tasty.HUnit

import           Test.Tasty.QuickCheck  as QC

vectors :: [ByteString]

vectors =

    [ ""

    , "The quick brown fox jumps over the lazy dog"

    , "The quick brown fox jumps over the lazy cog"

    , "abc"

    , "abcdbcdecdefdefgefghfghighijhijkijkljklmklmnlmnomnopnopq"

    , "abcdefghbcdefghicdefghijdefghijkefghijklfghijklmghijklmnhijklmnoijklmnopjklmnopqklmnopqrlmnopqrsmnopqrstnopqrstu"

    , B.replicate 1000000 0x61

    ]

answers :: [ByteString]

answers = map (B.filter (/= 0x20))

    [ "e3b0c442 98fc1c14 9afbf4c8 996fb924 27ae41e4 649b934c a495991b 7852b855"

    , "d7a8fbb3 07d78094 69ca9abc b0082e4f 8d5651e4 6d3cdb76 2d02d0bf 37c9e592"

    , "e4c4d8f3 bf76b692 de791a17 3e053211 50f7a345 b46484fe 427f6acc 7ecc81be"

    , "ba7816bf 8f01cfea 414140de 5dae2223 b00361a3 96177a9c b410ff61 f20015ad"

    , "248d6a61 d20638b8 e5c02693 0c3e6039 a33ce459 64ff2167 f6ecedd4 19db06c1"

    , "cf5b16a7 78af8380 036ce59e 7b049237 0b249b11 e8f07a51 afac4503 7afee9d1"

    , "cdc76e5c 9914fb92 81a1c7e2 84d73e67 f1809a48 a497200e 046d39cc c7112cd0"

    ]

ansXLTest :: ByteString

ansXLTest = B.filter (/= 0x20)

    "50e72a0e 26442fe2 552dc393 8ac58658 228c0cbf b1d2ca87 2ae43526 6fcd055e"

katTests :: [TestTree]

katTests

  | length vectors == length answers = map makeTest (zip3 [1::Int ..] vectors answers) ++ [xltest]

  | otherwise = error "vectors/answers length mismatch"

  where

    makeTest (i, v, r) = testGroup ("vec"++show i) $

        [ testCase "one-pass" (r @=? runTest v)

        , testCase "inc-1"    (r @=? runTestInc 1 v)

        , testCase "inc-2"    (r @=? runTestInc 2 v)

        , testCase "inc-3"    (r @=? runTestInc 3 v)

        , testCase "inc-4"    (r @=? runTestInc 4 v)

        , testCase "inc-5"    (r @=? runTestInc 5 v)

        , testCase "inc-7"    (r @=? runTestInc 7 v)

        , testCase "inc-8"    (r @=? runTestInc 8 v)

        , testCase "inc-9"    (r @=? runTestInc 9 v)

        , testCase "inc-16"   (r @=? runTestInc 16 v)

        , testCase "lazy-1"   (r @=? runTestLazy 1 v)

        , testCase "lazy-2"   (r @=? runTestLazy 2 v)

        , testCase "lazy-7"   (r @=? runTestLazy 7 v)

        , testCase "lazy-8"   (r @=? runTestLazy 8 v)

        , testCase "lazy-16"  (r @=? runTestLazy 16 v)

        ] ++

        [ testCase "lazy-63u"  (r @=? runTestLazyU 63 v) | B.length v > 63 ] ++

        [ testCase "lazy-65u"  (r @=? runTestLazyU 65 v) | B.length v > 65 ] ++

        [ testCase "lazy-97u"  (r @=? runTestLazyU 97 v) | B.length v > 97 ] ++

        [ testCase "lazy-131u" (r @=? runTestLazyU 131 v) | B.length v > 131 ]

    runTest :: ByteString -> ByteString

    runTest = B16.encode . IUT.hash

    runTestInc :: Int -> ByteString -> ByteString

    runTestInc i = B16.encode . IUT.finalize . myfoldl' IUT.update IUT.init . splitB i

    runTestLazy :: Int -> ByteString -> ByteString

    runTestLazy i = B16.encode . IUT.hashlazy . BL.fromChunks . splitB i

    -- force unaligned md5-blocks

    runTestLazyU :: Int -> ByteString -> ByteString

    runTestLazyU i = B16.encode . IUT.hashlazy . BL.fromChunks . map B.copy . splitB i

    ----

    xltest = testGroup "XL-vec"

        [ testCase "inc" (ansXLTest @=? (B16.encode . IUT.hashlazy) vecXL) ]

      where

        vecXL = BL.fromChunks (replicate 16777216 "abcdefghbcdefghicdefghijdefghijkefghijklfghijklmghijklmnhijklmno")

splitB :: Int -> ByteString -> [ByteString]

splitB l b

  | B.length b > l = b1 : splitB l b2

  | otherwise = [b]

  where

    (b1, b2) = B.splitAt l b

rfc4231Vectors :: [(ByteString,ByteString,ByteString)]

rfc4231Vectors = -- (secrect,msg,mac)

    [ (rep 20 0x0b, "Hi There", x"b0344c61d8db38535ca8afceaf0bf12b881dc200c9833da726e9376c2e32cff7")

    , ("Jefe", "what do ya want for nothing?", x"5bdcc146bf60754e6a042426089575c75a003f089d2739839dec58b964ec3843")

    , (rep 20 0xaa, rep 50 0xdd, x"773ea91e36800e46854db8ebd09181a72959098b3ef8c122d9635514ced565fe")

    , (B.pack [1..25], rep 50 0xcd, x"82558a389a443c0ea4cc819899f2083a85f0faa3e578f8077a2e3ff46729665b")

    , (rep 20 0x0c, "Test With Truncation", x"a3b6167473100ee06e0c796c2955552bfa6f7c0a6a8aef8b93f860aab0cd20c5")

    , (rep 131 0xaa, "Test Using Larger Than Block-Size Key - Hash Key First", x"60e431591ee0b67f0d8a26aacbf5b77f8e0bc6213728c5140546040f0ee37f54")

    , (rep 131 0xaa, "This is a test using a larger than block-size key and a larger than block-size data. The key needs to be hashed before being used by the HMAC algorithm.", x"9b09ffa71b942fcb27635fbcd5b0e944bfdc63644f0713938a7f51535c3a35e2")

    ]

  where

    x = fst.B16.decode

    rep n c = B.replicate n c

rfc4231Tests :: [TestTree]

rfc4231Tests = zipWith makeTest [1::Int ..] rfc4231Vectors

  where

    makeTest i (key, msg, mac) = testGroup ("vec"++show i) $

        [ testCase "hmac" (hex mac  @=? hex (IUT.hmac key msg))

        , testCase "hmaclazy" (hex mac  @=? hex (IUT.hmaclazy key lazymsg))

        ]

      where

        lazymsg = BL.fromChunks . splitB 1 $ msg

    hex = B16.encode

rfc5869Vectors :: [(Int,ByteString,ByteString,ByteString,ByteString)]

rfc5869Vectors = -- (l,ikm,salt,info,okm)

    [ (42, rep 22 0x0b, x"000102030405060708090a0b0c", x"f0f1f2f3f4f5f6f7f8f9", x"3cb25f25faacd57a90434f64d0362f2a2d2d0a90cf1a5a4c5db02d56ecc4c5bf34007208d5b887185865")

    , ( 82

      , x"000102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f202122232425262728292a2b2c2d2e2f303132333435363738393a3b3c3d3e3f404142434445464748494a4b4c4d4e4f"

      , x"606162636465666768696a6b6c6d6e6f707172737475767778797a7b7c7d7e7f808182838485868788898a8b8c8d8e8f909192939495969798999a9b9c9d9e9fa0a1a2a3a4a5a6a7a8a9aaabacadaeaf"

      , x"b0b1b2b3b4b5b6b7b8b9babbbcbdbebfc0c1c2c3c4c5c6c7c8c9cacbcccdcecfd0d1d2d3d4d5d6d7d8d9dadbdcdddedfe0e1e2e3e4e5e6e7e8e9eaebecedeeeff0f1f2f3f4f5f6f7f8f9fafbfcfdfeff"

      , x"b11e398dc80327a1c8e7f78c596a49344f012eda2d4efad8a050cc4c19afa97c59045a99cac7827271cb41c65e590e09da3275600c2f09b8367793a9aca3db71cc30c58179ec3e87c14c01d5c1f3434f1d87"

      )

    , ( 42, rep 22 0x0b, "", "", x"8da4e775a563c18f715f802a063c5a31b8a11f5c5ee1879ec3454e5f3c738d2d9d201395faa4b61a96c8")

    ]

  where

    x = fst.B16.decode

    rep n c = B.replicate n c

rfc5869Tests :: [TestTree]

rfc5869Tests = zipWith makeTest [1::Int ..] rfc5869Vectors

  where

    makeTest i (l,ikm,salt,info,okm) = testGroup ("vec"++show i) $

        [ testCase "hkdf" (hex okm @=? hex (IUT.hkdf ikm salt info l)) ]

    hex = B16.encode

-- define own 'foldl' here to avoid RULE rewriting to 'hashlazy'

myfoldl' :: (b -> a -> b) -> b -> [a] -> b

myfoldl' f z0 xs0 = lgo z0 xs0

  where

    lgo z []     = z

    lgo z (x:xs) = let z' = f z x

                   in z' `seq` lgo z' xs

newtype RandBS = RandBS { unRandBS :: ByteString }

newtype RandLBS = RandLBS BL.ByteString

instance Arbitrary RandBS where

    arbitrary = fmap (RandBS . B.pack) arbitrary

    shrink (RandBS x) = fmap RandBS (go x)

      where

        go bs = zipWith B.append (B.inits bs) (tail $ B.tails bs)

instance Show RandBS where

    show (RandBS x) = "RandBS {len=" ++ show (B.length x)++"}"

instance Arbitrary RandLBS where

    arbitrary = fmap (RandLBS . BL.fromChunks . map unRandBS) arbitrary

instance Show RandLBS where

    show (RandLBS x) = "RandLBS {len=" ++ show (BL.length x) ++ ", chunks=" ++ show (length $ BL.toChunks x)++"}"

refImplTests :: [TestTree]

refImplTests =

    [ testProperty "hash" prop_hash

    , testProperty "hashlazy" prop_hashlazy

    , testProperty "hashlazyAndLength" prop_hashlazyAndLength

    , testProperty "hmac" prop_hmac

    , testProperty "hmaclazy" prop_hmaclazy

    , testProperty "hmaclazyAndLength" prop_hmaclazyAndLength

    ]

  where

    prop_hash (RandBS bs)

        = ref_hash bs == IUT.hash bs

    prop_hashlazy (RandLBS bs)

        = ref_hashlazy bs == IUT.hashlazy bs

    prop_hashlazyAndLength (RandLBS bs)

        = ref_hashlazyAndLength bs == IUT.hashlazyAndLength bs

    prop_hmac (RandBS k) (RandBS bs)

        = ref_hmac k bs == IUT.hmac k bs

    prop_hmaclazy (RandBS k) (RandLBS bs)

        = ref_hmaclazy k bs == IUT.hmaclazy k bs

    prop_hmaclazyAndLength (RandBS k) (RandLBS bs)

        = ref_hmaclazyAndLength k bs == IUT.hmaclazyAndLength k bs

    ref_hash :: ByteString -> ByteString

    ref_hash = ref_hashlazy . fromStrict

    ref_hashlazy :: BL.ByteString -> ByteString

    ref_hashlazy = toStrict . REF.bytestringDigest . REF.sha256

    ref_hashlazyAndLength :: BL.ByteString -> (ByteString,Word64)

    ref_hashlazyAndLength x = (ref_hashlazy x, fromIntegral (BL.length x))

    ref_hmac :: ByteString -> ByteString -> ByteString

    ref_hmac secret = ref_hmaclazy secret . fromStrict

    ref_hmaclazy :: ByteString -> BL.ByteString -> ByteString

    ref_hmaclazy secret = toStrict . REF.bytestringDigest . REF.hmacSha256 (fromStrict secret)

    ref_hmaclazyAndLength :: ByteString -> BL.ByteString -> (ByteString,Word64)

    ref_hmaclazyAndLength secret msg = (ref_hmaclazy secret msg, fromIntegral (BL.length msg))

    -- toStrict/fromStrict only available with bytestring-0.10 and later

    toStrict = B.concat . BL.toChunks

    fromStrict = BL.fromChunks . (:[])

main :: IO ()

main = defaultMain $ testGroup "cryptohash-sha256"

    [ testGroup "KATs" katTests

    , testGroup "RFC4231" rfc4231Tests

    , testGroup "RFC5869" rfc5869Tests

    , testGroup "REF" refImplTests

    ]