{-# LANGUAGE CPP                #-}
{-# LANGUAGE DeriveGeneric      #-}
{-# LANGUAGE OverloadedStrings  #-}
{-# LANGUAGE RankNTypes         #-}
{-# LANGUAGE RecordWildCards    #-}
{-# LANGUAGE TemplateHaskell    #-}
{-# LANGUAGE QuasiQuotes        #-}
{-# LANGUAGE DeriveDataTypeable #-}

module Utils (
    CBCMode
  , CFBMode
  , CTRMode

  , propRoundTrip
  , genPropRoundTrip
  , groupRoundTrip

  , modifyId
  , modifyBase64
  , modifyCookie
  , modifyPayload
  , modifyMAC
  , modifyExpiration

  , checkEquals
  , checkException
  , checkSessionDeserializationFailed
  , checkIncorrectMAC
  , checkCookieExpired
  ) where

import Control.Monad.Catch (MonadThrow (..))
import Crypto.Cipher.Types (BlockCipher(..), Cipher(..))
import Crypto.Hash (HashAlgorithm, SHA256(..), SHA384(..), SHA512(..))
import Crypto.Random (drgNew)
import Data.ByteString (ByteString)
import Data.Default (def)
import Data.List (intercalate)
import Data.Proxy (Proxy(..))
import Data.Serialize (Serialize)
import Data.Tagged (Tagged(..), unTagged)
import Data.Time (addUTCTime)
import Data.Typeable (Typeable, typeRep)
import GHC.Generics (Generic)
import Language.Haskell.TH.Syntax (Name, Type(..), Exp(..), Q)
import Servant.Server.Experimental.Auth.Cookie
import Test.Hspec (Spec, Selector, Expectation, shouldThrow, shouldReturn)
import Test.Hspec.QuickCheck (prop)
import Test.QuickCheck (Arbitrary(..), Gen, vectorOf, oneof, choose, sized)
import qualified Data.ByteString as BS
import qualified Data.ByteString.Char8 as BSC8

#if !MIN_VERSION_base(4,8,0)
import Control.Applicative ((<$>), (<*>))
#endif


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

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))]

type Updater a = a -> IO a

type Modifier c a
  =  forall k. (ServerKeySet k)
  => AuthCookieSettings
  -> RandomSource
  -> k
  -> Proxy a
  -> Updater c

type EncryptedCookieModifier a = Modifier (Tagged SerializedEncryptedCookie ByteString) a


insideBase64 :: Modifier (Tagged EncryptedCookie ByteString) a -> EncryptedCookieModifier a
insideBase64 f = \acs rs sks p x -> base64Decode x >>= fmap base64Encode . f acs rs sks p

insideSerializedCookie :: Modifier Cookie a -> EncryptedCookieModifier a
insideSerializedCookie f = insideBase64 $ \acs rs sks p x -> cerealDecode x >>= fmap cerealEncode . f acs rs sks p

insideConsistentCookie :: Modifier Cookie a -> EncryptedCookieModifier a
insideConsistentCookie f = insideSerializedCookie $ \(acs@AuthCookieSettings {..}) rs sks p c -> do
  sk <- (Tagged . fst) <$> getKeys sks
  cookiePayload' <- cookiePayload <$> f acs rs sks p c
  cookiePadding' <- mkPadding rs acsCipher cookiePayload'
  let c' = c {
          cookiePayload = cookiePayload'
        , cookiePadding = cookiePadding'
        }
  return c' { cookieMAC = mkMAC acsHashAlgorithm sk c' }

insideEncryptedCookie :: Modifier (Tagged PayloadBytes ByteString) a -> EncryptedCookieModifier a
insideEncryptedCookie f = insideConsistentCookie $ \(acs@AuthCookieSettings {..}) rs sks p (c@Cookie {..}) -> do
  sk <- (Tagged . fst) <$> getKeys sks
  key <- mkCookieKey acsCipher acsHashAlgorithm sk cookieIV
  cookiePayload' <- applyCipherAlgorithm acsDecryptAlgorithm cookieIV key cookiePayload
                >>= f acs rs sks p
                >>= applyCipherAlgorithm acsEncryptAlgorithm cookieIV key
  return c { cookiePayload = cookiePayload' }

nullify :: Tagged a ByteString -> IO (Tagged a ByteString)
nullify = return . Tagged . const BS.empty . unTagged

updatePayload :: (Tagged PayloadBytes ByteString -> IO (Tagged PayloadBytes ByteString)) -> Cookie -> IO Cookie
updatePayload f c = (f $ cookiePayload c) >>= \cookiePayload' -> return c { cookiePayload = cookiePayload' }

cerealDecode' :: (Serialize a, MonadThrow m) => Proxy a -> Tagged b ByteString -> m (PayloadWrapper a)
cerealDecode' _ = cerealDecode

cerealEncode' :: (Serialize a) => Proxy a -> PayloadWrapper a -> Tagged b ByteString
cerealEncode' _ = cerealEncode


modifyId :: EncryptedCookieModifier a
modifyId _ _ _ _ = return . id

modifyBase64 :: EncryptedCookieModifier a
modifyBase64 _ _ _ _ = return . fmap (BSC8.scanl1 (\c c' -> if c == '_' then c' else '_'))

modifyCookie :: EncryptedCookieModifier a
modifyCookie = insideBase64 $ \_ _ _ _ -> nullify

modifyPayload :: EncryptedCookieModifier a
modifyPayload = insideConsistentCookie $ \_ _ _ _ -> updatePayload nullify

modifyMAC :: EncryptedCookieModifier a
modifyMAC = insideSerializedCookie $ \_ _ _ _ -> updatePayload nullify

modifyExpiration :: (Serialize a) => EncryptedCookieModifier a
modifyExpiration = insideEncryptedCookie $ \AuthCookieSettings {..} _ _ p s -> do
  r <- cerealDecode' p s
  return $ cerealEncode' p r { pwExpiration = addUTCTime (-acsMaxAge * 2) (pwExpiration r) }


type SessionChecker a = (Show a, Eq a) => a -> IO a -> Expectation

checkEquals :: SessionChecker a
checkEquals = flip shouldReturn

checkException :: Selector AuthCookieException -> SessionChecker a
checkException e = \_ -> flip shouldThrow e

checkSessionDeserializationFailed :: SessionChecker a
checkSessionDeserializationFailed = checkException sel where
  sel :: AuthCookieException -> Bool
  sel (SessionDeserializationFailed _) = True
  sel _                                = False

checkIncorrectMAC :: SessionChecker a
checkIncorrectMAC = checkException sel where
  sel :: AuthCookieException -> Bool
  sel (IncorrectMAC _) = True
  sel _                = False

checkCookieExpired :: SessionChecker a
checkCookieExpired = checkException sel where
  sel :: AuthCookieException -> Bool
  sel (CookieExpired _ _) = True
  sel _                   = False


roundTrip
  :: (Serialize a)
  => AuthCookieSettings
  -> EncryptedCookieModifier a
  -> Proxy a
  -> a
  -> IO a
roundTrip settings modify p x = do
  rs <- mkRandomSource drgNew 1000
  sk <- mkPersistentServerKey <$> generateRandomBytes 16
  encryptSession settings rs sk def x >>= modify settings rs sk p >>= (fmap epwSession . decryptSession settings sk)


class (HashAlgorithm h) => NamedHashAlgorithm h where
  hashName :: h -> String

instance NamedHashAlgorithm SHA512 where
  hashName _ = show SHA512

instance NamedHashAlgorithm SHA384 where
  hashName _ = show SHA384

instance NamedHashAlgorithm SHA256 where
  hashName _ = show SHA256


class BlockCipherMode m where
  modeName    :: m -> String
  modeEncrypt :: (BlockCipher c) => m -> CipherAlgorithm c
  modeDecrypt :: (BlockCipher c) => m -> CipherAlgorithm c

data CBCMode
instance BlockCipherMode CBCMode where
  modeName    _ = "CBC"
  modeEncrypt _ = cbcEncrypt
  modeDecrypt _ = cbcDecrypt

data CFBMode
instance BlockCipherMode CFBMode where
  modeName    _ = "CFB"
  modeEncrypt _ = cfbEncrypt
  modeDecrypt _ = cfbDecrypt

data CTRMode
instance BlockCipherMode CTRMode where
  modeName    _ = "CTR"
  modeEncrypt _ = ctrCombine
  modeDecrypt _ = ctrCombine


mkSettings
  :: ( NamedHashAlgorithm h
     , BlockCipher c
     , BlockCipherMode m
  ) => Proxy h
    -> Proxy c
    -> Proxy m
    -> AuthCookieSettings
mkSettings h c m = (def $) $ \(AuthCookieSettings{..}) -> AuthCookieSettings {
    acsHashAlgorithm    = h
  , acsCipher           = c
  , acsEncryptAlgorithm = modeEncrypt (unProxy m)
  , acsDecryptAlgorithm = modeDecrypt (unProxy m)
  , ..}

mkTestName
  :: ( NamedHashAlgorithm h
     , BlockCipher c
     , Serialize a, Arbitrary a, Show a, Eq a, Typeable a
     , BlockCipherMode m
  ) => Proxy h
    -> Proxy c
    -> Proxy m
    -> Proxy a
    -> String
mkTestName h c m a = intercalate "_" [
    hashName   $ unProxy h
  , cipherName $ unProxy c
  , modeName   $ unProxy m
  ] ++ (" (" ++ (show . typeRep $ a) ++ ")")

propRoundTrip
  :: ( NamedHashAlgorithm h
     , BlockCipher c
     , Serialize a, Arbitrary a, Show a, Eq a, Typeable a
     , BlockCipherMode m
  ) => Proxy h
    -> Proxy c
    -> Proxy m
    -> Proxy a
    -> EncryptedCookieModifier a
    -> SessionChecker a
    -> Spec
propRoundTrip h c m a modify check = prop (mkTestName h c m a) $
  \x -> check x (roundTrip (mkSettings h c m) modify a x)


mkProxy :: Type -> Q Exp
mkProxy t = [| Proxy :: Proxy $(return t) |]

genPropRoundTrip
  :: Name  -- ^ Hash name
  -> Name  -- ^ Cipher name
  -> Name  -- ^ Block cipher mode
  -> Name  -- ^ Session type name
  -> Name  -- ^ Modifier name
  -> Name  -- ^ Checker name
  -> Q Exp -- ^ Function of type (EncryptedCookieModifier -> Spec)
genPropRoundTrip h c m a modify check = [|
  propRoundTrip
    $(mkProxy $ PromotedT h)
    $(mkProxy $ PromotedT c)
    $(mkProxy $ PromotedT m)
    $(mkProxy $ (PromotedT ''Tree) `AppT` (PromotedT a))
    $(return $ VarE modify)
    $(return $ VarE check)
  |]

groupRoundTrip :: [Q Exp] -> Q Exp
groupRoundTrip qs = [| sequence_ $(ListE <$> sequence qs) |]