{-# LANGUAGE CPP               #-}
{-# LANGUAGE DeriveGeneric     #-}
{-# LANGUAGE FlexibleContexts  #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RankNTypes        #-}
{-# LANGUAGE RecordWildCards   #-}

module Main (main) where

import           Control.Concurrent                      (threadDelay)
import           Control.Monad.IO.Class                  (MonadIO, liftIO)
import           Crypto.Cipher.AES                       (AES128, AES192,
                                                          AES256)
import           Crypto.Cipher.Types
import           Crypto.Hash                             (HashAlgorithm, SHA256,
                                                          SHA384, SHA512)
import           Crypto.Random                           (drgNew)
import           Data.ByteString                         (ByteString)
import qualified Data.ByteString                         as BS
import           Data.Default
import           Data.Proxy
import           Data.Serialize                          (Serialize)
import           Data.Time
import           GHC.Generics                            (Generic)
import           Servant.Server.Experimental.Auth.Cookie
import           Test.Hspec
import           Test.QuickCheck

#if !MIN_VERSION_base(4,8,0)
import           Control.Applicative
#endif

main :: IO ()
main = hspec spec

spec :: Spec
spec = do
  describe "RandomSource"        randomSourceSpec
  describe "PersistentServerKey" persistentServerKeySpec
  describe "RenewalKeySet"       renewalKeySetSpec
  describe "Cookie"              cookieSpec
  describe "Session"             sessionSpec

randomSourceSpec :: Spec
randomSourceSpec = do
  context "when getRandomBytes is called consequently not exceeding threshould" $
    it "does not produces the same result" $ do
      rs <- mkRandomSource drgNew 100
      s1 <- getRandomBytes rs 10
      s2 <- getRandomBytes rs 10
      s1 `shouldNotBe` s2
  context "when two sources are created with the same parameters" $
    it "they do not produce the same output" $ do
      rs1 <- mkRandomSource drgNew 100
      rs2 <- mkRandomSource drgNew 100
      s1  <- getRandomBytes rs1 10
      s2  <- getRandomBytes rs2 10
      s1 `shouldNotBe` s2
  context "after resetting" $
    it "does not produce the same result" $ do
      rs <- mkRandomSource drgNew 10
      s1 <- getRandomBytes rs 10
      s2 <- getRandomBytes rs 10
      s1 `shouldNotBe` s2

persistentServerKeySpec :: Spec
persistentServerKeySpec = do
  context "when creating a new persistent server key" $ do
    let keySize = 64
    let getSK = mkPersistentServerKey <$> generateRandomBytes keySize
            >>= getKeys

    it "has correct size" $ do
      (k, _) <- getSK
      BS.length k `shouldNotBe` (keySize `div` 8)

    it "has no rotated keys" $ do
      (_, ks) <- getSK
      length ks `shouldBe` 0


renewalKeySetSpec :: Spec
renewalKeySetSpec = spec' where

  keySize :: Int
  keySize = 16

  rkshNewState :: (MonadIO m)
    => NominalDiffTime
    -> ([ServerKey], UTCTime)
    -> m ([ServerKey], UTCTime)
  rkshNewState _ (keys, _) = liftIO $ (,)
    <$> (fmap (:keys) $ generateRandomBytes keySize)
    <*> getCurrentTime

  rkshNeedUpdate :: (MonadIO m)
    => NominalDiffTime
    -> ([ServerKey], UTCTime)
    -> m Bool
  rkshNeedUpdate dt (_, t) = liftIO $ getCurrentTime >>= return . ((dt `addUTCTime` t) <)

  rkshRemoveKey :: (MonadIO m)
    => NominalDiffTime
    -> ServerKey
    -> m ()
  rkshRemoveKey _ _ = return ()

  spec' = do
    let makeSK = mkRenewableKeySet
          RenewableKeySetHooks {..}
          (fromIntegral (1 :: Integer))
          (UTCTime (toEnum 0) 0)

    context "when accessing a renewable key set" $ do
      it "updates the keys when needed" $ do
        sk <- makeSK

        (k, ks) <- getKeys sk
        BS.length k `shouldNotBe` 0

        (_, ks') <- threadDelay 1500000 >> getKeys sk
        ks' `shouldBe` (k:ks)

      it "doesn't update the keys when not needed" $ do
        sk <- makeSK
        k  <- getKeys sk
        k' <- getKeys sk
        k' `shouldBe` k

    context "when removing key from a renewable key set" $ do
      it "removes specified key" $ do
        sk <- makeSK
        (_, ks) <- getKeys sk >> threadDelay 1500000 >> getKeys sk
        length ks `shouldNotBe` 0

        let k = head ks
        removeKey sk k
        (_, ks') <- getKeys sk
        (k:ks') `shouldBe` ks


cookieSpec :: Spec
cookieSpec = do
  context "when used with different encryption/decryption algorithms" $ do
    it "works in CBC mode" $
      testCipher (Proxy :: Proxy SHA256) (Proxy :: Proxy AES256) cbcEncrypt cbcDecrypt 64
    it "works in CFB mode" $
      testCipher (Proxy :: Proxy SHA256) (Proxy :: Proxy AES256) cfbEncrypt cfbDecrypt 64
    it "works in CTR combine mode" $
      testCipher (Proxy :: Proxy SHA256) (Proxy :: Proxy AES256) ctrCombine ctrCombine 100
  context "when used with different ciphers" $ do
    it "works with AES256" $
      testCipher (Proxy :: Proxy SHA256) (Proxy :: Proxy AES256) ctrCombine ctrCombine 100
    it "works with AES192" $
      testCipher (Proxy :: Proxy SHA256) (Proxy :: Proxy AES192) ctrCombine ctrCombine 100
    it "works with AES128" $
      testCipher (Proxy :: Proxy SHA256) (Proxy :: Proxy AES128) ctrCombine ctrCombine 100
  context "when used with different hash algorithms" $ do
    it "works with SHA512" $
      testCipher (Proxy :: Proxy SHA512) (Proxy :: Proxy AES256) ctrCombine ctrCombine 100
    it "works with SHA384" $
      testCipher (Proxy :: Proxy SHA384) (Proxy :: Proxy AES256) ctrCombine ctrCombine 100
    it "works with SHA256" $
      testCipher (Proxy :: Proxy SHA256) (Proxy :: Proxy AES256) ctrCombine ctrCombine 100
  context "when cookie is corrupted" $
    it "throws" $
      let selectIncorrectMAC (IncorrectMAC _) = True
          selectIncorrectMAC _                = False
      in testCustomCookie
        (mkCookie 10 100)
        (BS.drop 1)
        selectIncorrectMAC
  context "when cookie has expired" $
    it "throws CookieExpired" $
      let selectCookieExpired (CookieExpired _ _) = True
          selectCookieExpired _                   = False
      in testCustomCookie
        (mkCookie 0 100)
        id
        selectCookieExpired

testCipher :: (HashAlgorithm h, BlockCipher c)
  => Proxy h           -- ^ Hash algorithm
  -> Proxy c           -- ^ Cipher
  -> CipherAlgorithm c -- ^ Encryption algorithm
  -> CipherAlgorithm c -- ^ Decryption algorithm
  -> Int               -- ^ Payload size
  -> Expectation
testCipher h c encryptAlgorithm decryptAlgorithm size = do
  cookie <- mkCookie 10 size
  result <- cipherId h c encryptAlgorithm decryptAlgorithm cookie id
  cookieIV             result `shouldBe` cookieIV             cookie
  diffUTCTime (cookieExpirationTime cookie) (cookieExpirationTime result)
    `shouldSatisfy` (< fromIntegral (1 :: Integer))
  cookiePayload        result `shouldBe` cookiePayload        cookie

testCustomCookie
  :: IO Cookie
  -> (ByteString -> ByteString)
  -> Selector AuthCookieException
  -> Expectation
testCustomCookie mkCookie' encryptionHook selector = do
  cookie <- mkCookie'
  cipherId
    (Proxy :: Proxy SHA256)
    (Proxy :: Proxy AES256)
    ctrCombine ctrCombine
    cookie
    encryptionHook
    `shouldThrow` selector

mkCookie :: Int -> Int -> IO Cookie
mkCookie dt size = do
  rs         <- mkRandomSource drgNew 1000
  iv         <- getRandomBytes rs 16
  expiration <- addUTCTime (fromIntegral dt) <$> getCurrentTime
  payload    <- getRandomBytes rs size
  return Cookie
    { cookieIV             = iv
    , cookieExpirationTime = expiration
    , cookiePayload        = payload }

cipherId :: (HashAlgorithm h, BlockCipher c)
  => Proxy h           -- ^ Hash algorithm
  -> Proxy c           -- ^ Cipher
  -> CipherAlgorithm c -- ^ Encryption algorithm
  -> CipherAlgorithm c -- ^ Decryption algorithm
  -> Cookie            -- ^ 'Cookie' to encrypt
  -> (BS.ByteString -> BS.ByteString) -- ^ Encryption hook
  -> IO Cookie         -- ^ Restored 'Cookie'
cipherId h c encryptAlgorithm decryptAlgorithm cookie encryptionHook = do
  sk <- mkPersistentServerKey <$> generateRandomBytes 16

  let sts =
        case def of
          AuthCookieSettings {..} -> AuthCookieSettings
            { acsEncryptAlgorithm = encryptAlgorithm
            , acsDecryptAlgorithm = decryptAlgorithm
            , acsHashAlgorithm    = h
            , acsCipher           = c
            , .. }
  encryptCookie sts sk cookie >>= (fmap wmData . decryptCookie sts sk . fmap encryptionHook)

sessionSpec :: Spec
sessionSpec = do
  let treesOfInt = Proxy :: Proxy Int
      treesOfString = Proxy :: Proxy String

  context "when session is encrypted and decrypted" $ do
    it "is not distorted in any way (1)" $
      property $ \session -> encryptThenDecrypt treesOfInt def session `shouldReturn` session
    it "is not distored in any way (2)" $
      property $ \session -> encryptThenDecrypt treesOfString def session `shouldReturn` session
  context "when session is encrypted and decrypted (CBC mode)" $ do
    let sts = case def of
                AuthCookieSettings{..} -> AuthCookieSettings
                                          { acsEncryptAlgorithm = cbcEncrypt
                                          , acsDecryptAlgorithm = cbcDecrypt
                                          , .. }
    it "is not distorted in any way (1)" $
      property $ \session -> encryptThenDecrypt treesOfInt sts session `shouldReturn` session
    it "is not distored in any way (2)" $
      property $ \session -> encryptThenDecrypt treesOfString sts session `shouldReturn` session
  context "when session is encrypted and decrypted (CFB mode)" $ do
    let sts =
          case def of
            AuthCookieSettings {..} -> AuthCookieSettings
                                       { acsEncryptAlgorithm = cfbEncrypt
                                       , acsDecryptAlgorithm = cfbDecrypt
                                       , .. }
    it "is not distorted in any way (1)" $
      property $ \session -> encryptThenDecrypt treesOfInt sts session `shouldReturn` session
    it "is not distored in any way (2)" $
      property $ \session -> encryptThenDecrypt treesOfString sts session `shouldReturn` session

encryptThenDecrypt :: Serialize a
  => Proxy a
  -> AuthCookieSettings
  -> Tree a
  -> IO (Tree a)
encryptThenDecrypt _ settings x = do
  rs <- mkRandomSource drgNew 1000
  sk <- mkPersistentServerKey <$> generateRandomBytes 16
  encryptSession settings rs sk x >>= (fmap wmData . decryptSession settings sk)

data Tree a = Leaf a | Node a [Tree a] deriving (Eq, Show, Generic)

instance Serialize a => Serialize (Tree a)

instance Arbitrary a => Arbitrary (Tree a) where
  arbitrary = sized arbitraryTree

arbitraryTree :: Arbitrary a => Int -> Gen (Tree a)
arbitraryTree 0 = Leaf <$> arbitrary
arbitraryTree n = do
  l <- choose (0, n `quot` 2)
  oneof
    [ Leaf <$> arbitrary
    , Node <$> arbitrary <*> vectorOf l (arbitraryTree (n `quot` 2))]