{-# LANGUAGE Haskell2010               #-}

{-# LANGUAGE DeriveFunctor             #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE OverloadedStrings         #-}
{-# LANGUAGE Rank2Types                #-}
{-# LANGUAGE ScopedTypeVariables       #-}

-- SPDX-License-Identifier: GPL-2.0-or-later

module Main (main) where

import           Control.Exception     (catchJust)
import           Control.Monad
import           Data.ByteString       (ByteString)
import qualified Data.ByteString       as BS
import qualified Data.ByteString.Lazy  as BL
import           Data.Maybe
import           System.IO             (hPutStrLn, stderr)
import           System.IO.Error       (isDoesNotExistError)
import           Test.Tasty
import           Test.Tasty.HUnit

import           CAVS

import           Crypto.Cipher.AES.GCM as GCM

-- WARNING: the code that follows will make you cry;
--          a safety pig is provided below for your benefit.
--
--                          _
--  _._ _..._ .-',     _.._(`))
-- '-. `     '  /-._.-'    ',/
--    )         \            '.
--   / _    _    |             \
--  |  a    a    /              |
--  \   .-.                     ;
--   '-('' ).-'       ,'       ;
--      '-;           |      .'
--         \           \    /
--         | 7  .__  _.-\   \
--         | |  |  ``/  /`  /
--        /,_|  |   /,_/   /
--           /,_/      '`-'
--

main :: IO ()
main = do
    rspGrps <- forM rspFiles $ \(opt,fn) -> do
      let fn' = "testdata/" ++ fn ++ ".rsp"
      tvs <- catchJust (\e -> if isDoesNotExistError e && opt then Just () else Nothing)
                       (parseTestVecs <$> BL.readFile fn')
                       (\() -> hPutStrLn stderr ("optional " ++ fn ++ " test vectors not found; see testdata/README for how to procure those") >> pure [])

      pure $ testGroup fn $ map (uncurry mkTestGrp) tvs

    defaultMain $ testGroup "aes-gcm"
        [ testGroup "large data" largeTests
        , testGroup "test vectors" rspGrps
        ]
  where
    rspFiles = [ (False,"gcm-spec")
               , (True,"gcmDecrypt128")
               , (True,"gcmDecrypt192")
               , (True,"gcmDecrypt256")
               , (True,"gcmEncryptExtIV128")
               , (True,"gcmEncryptExtIV192")
               , (True,"gcmEncryptExtIV256")
               ]

    mkTestGrp :: TestVecHdr -> [TestVec] -> TestTree
    mkTestGrp tvh tvs = testGroup glabel $ concatMap mkTestCase tvs
      where
        glabel = mconcat [ "AES-", show (tv'Keylen tvh), "-GCM"
                         , " iv=", show (tv'IVlen tvh)
                         , " pt=", show (tv'PTlen tvh)
                         , " aad=", show (tv'AADlen tvh)
                         , " tag=", show (tv'Taglen tvh)
                        ]

-- aes-256-gcm with Key=0^256 IV=0^96 and max-size PT has T=0x646f81304169ebaff70a564a6edf79ba
largeTests :: [TestTree]
largeTests =
  [ testCase ("AES-" ++ show (8*klen) ++ "-GCM 1GiB") $ do
      let key  = fromJust $ someKey'fromByteString (BS.replicate klen 0x00)
          iv96 = fromJust $ iv'fromByteString (BS.replicate 12 0x00)

      (ectx,dctx) <- case key of
        SomeKey k -> (,) <$> encryptInit k iv96
                         <*> decryptInit k iv96

      -- let pt'ref = BS.replicate (1024*1024) 0xa5 -- 1 MiB
      let pt'in = BS.replicate (1024*1024) 0x00 -- 1 MiB

      -- 1 GiB of AAD
      forM_ [1 .. 1024 :: Int] $ \_ -> encryptUpdateAAD ectx pt'in
      forM_ [1 .. 1024 :: Int] $ \_ -> decryptUpdateAAD dctx pt'in

      -- 1 GiB of PT/CT
      forM_ [1 .. 1024] $ \j -> do
        ct <- encryptUpdate ectx pt'in
        pt <- decryptUpdate dctx ct
        assertEqual ("decrypt(encrypt(chunk)) != chunk #" ++ show (j :: Int))  pt'in pt

      tag <- encryptFinalize ectx

      valid <- decryptFinalize dctx (tag :: Tag Bits'128)

      assertBool "auth tag mismatched" valid

      assertEqual "auth tag" tag'ref (hex $ tag'toByteString tag)

      pure ()
  | (klen,tag'ref) <- [(16,"c15882df03ccc44a1ade08a947478855")
                      ,(24,"3dbf18efe8699ececd1a571c9737dcc3")
                      ,(32,"95cb9913f41ca00ec7c6d3ad6d5c7ede")
                      ]
  ]


tv'SomeKey :: TestVec -> SomeKey
tv'SomeKey tv = fromMaybe (error "invalid key-length in test-case") $ someKey'fromByteString (tv'Key tv)

tv'SomeTag :: TestVec -> SomeTag
tv'SomeTag tv = fromMaybe (error "invalid tag-length in test-case") $ someTag'fromByteString (tv'Tag tv)

mkTestCase :: TestVec -> [TestTree]
mkTestCase tv = concat $
    [ [ testCase ("encrypt           #"++show(tv'Count tv)) $
          Right (ct,tag) @=? encryptViaSome key iv aad pt tagLen

      | Just pt <- [mpt]
      ]
    , [ testCase ("encrypt...        #"++show(tv'Count tv)) $ do
          ctx <- case key of SomeKey k -> encryptInit k iv
          unless (BS.null aad) $ encryptUpdateAAD ctx aad
          ct'new <- encryptUpdate ctx pt
          ct @=? ct'new

          tag'new <- fromJust $ someTagFunctor tagLen (encryptFinalize ctx)
          tag @=? tag'new

      | Just pt <- [mpt]
      ]
    , [ testCase ("decrypt    ("++lab++") #"++show(tv'Count tv)) $
          Right mpt @=? decryptViaSome key iv tag aad ct
      ]
    , [ testCase ("decrypt... ("++lab++") #"++show(tv'Count tv)) $ do
          ctx <- case key of SomeKey k -> decryptInit k iv
          unless (BS.null aad) $ decryptUpdateAAD ctx aad
          pt'new <- decryptUpdate ctx ct

          valid <- case tag of SomeTag t -> decryptFinalize ctx t
          case mpt of
            Nothing -> assertBool ("unexpected auth tag match") (not valid)
            Just pt -> do
              assertBool ("unexpected auth tag mismatch") valid
              pt @=? pt'new
      ]
    ]
  where
    aad = tv'AAD tv
    key = tv'SomeKey tv
    tag = tv'SomeTag tv
    tagLen = BS.length (tv'Tag tv)
    iv = fromMaybe (error "invalid iv-length in test-case") (iv'fromByteString $ tv'IV tv)
    ct = tv'CT tv

    (mpt,lab) = case tv'PT tv of
      PT _ pt -> (Just pt,"pass")
      P'FAIL  -> (Nothing,"fail")

----------------------------------------------------------------------------
-- SomeKey/SomeTag based wrappers

decryptViaSome :: SomeKey -> IV -> SomeTag -> ByteString -> ByteString -> Either AesGcmError (Maybe ByteString)
decryptViaSome (SomeKey key) iv (SomeTag tag) aad pt = decrypt key iv tag aad pt

encryptViaSome :: SomeKey -> IV -> ByteString -> ByteString -> Int -> Either AesGcmError (ByteString, SomeTag)
encryptViaSome (SomeKey key) iv aad pt tlen
    = unERes $ fromMaybe (error "invalid tag-length in test-case") $ someTagFunctor tlen $ ERes $
      encrypt key iv aad pt

-- type helper
newtype EncRes a = ERes { unERes :: (Either AesGcmError (ByteString, a)) } deriving Functor

{-# INLINE someTagFunctor  #-}
someTagFunctor :: forall f . Functor f => Int -> (forall tlen . (KnownTagLength tlen) => f (Tag tlen)) -> Maybe (f SomeTag)
someTagFunctor tlen ftag = case tlen of
   4  -> Just $ fmap SomeTag (ftag :: f (Tag Bits'32))
   8  -> Just $ fmap SomeTag (ftag :: f (Tag Bits'64))
   12 -> Just $ fmap SomeTag (ftag :: f (Tag Bits'96))
   13 -> Just $ fmap SomeTag (ftag :: f (Tag Bits'104))
   14 -> Just $ fmap SomeTag (ftag :: f (Tag Bits'112))
   15 -> Just $ fmap SomeTag (ftag :: f (Tag Bits'120))
   16 -> Just $ fmap SomeTag (ftag :: f (Tag Bits'128))
   _  -> Nothing