{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE TupleSections #-}

import AuthAPI (app, authSettings, LoginForm(..), homePage, loginPage, Account(..))
import Control.Monad.IO.Class (liftIO)
import Crypto.Random (drgNew)
import Data.Default (def)
import Data.Int (Int64)
import Data.Maybe (fromMaybe)
import Data.Time.Clock (UTCTime(..))
import Network.HTTP.Media.RenderHeader (renderHeader)
import Network.HTTP.Types (Header, methodGet, methodPost, hContentType, hCookie)
import Network.Wai.Test (SResponse(..))
import Prelude ()
import Prelude.Compat
import Servant (FormUrlEncoded, contentType)
import Servant (Proxy(..))
import Servant.Server.Experimental.Auth.Cookie
import Test.Hspec (Spec, hspec, describe, context, it, shouldBe)
import Test.Hspec.Wai (WaiSession, WaiExpectation, shouldRespondWith, with, request, get)
import Text.Blaze (Markup)
import Text.Blaze.Renderer.Utf8 (renderMarkup)
import Web.Cookie (parseCookies)
import qualified Data.ByteString as BS
import qualified Data.ByteString.Lazy as BSL
import qualified Data.ByteString.Lazy.Char8 as BSLC8

#if MIN_VERSION_hspec_wai (0,7,0)
import Test.Hspec.Wai.Matcher (bodyEquals, ResponseMatcher(..), MatchBody(..))
#else
import Test.Hspec.Wai (matchBody)
#endif

#if MIN_VERSION_servant (0,9,0)
import Web.FormUrlEncoded (ToForm, toForm, urlEncodeForm)
#else
import Servant (ToFormUrlEncoded, mimeRender)
#endif

#if MIN_VERSION_servant (0,9,1) && MIN_VERSION_directory (1,2,5)
import FileKeySet (mkFileKeySet, FileKSParams(..), mkFileKey)
import Control.Arrow ((***))
import Control.Monad (void, when)
import Data.Monoid ((<>))
import Control.Exception.Base (bracket)
import Network.HTTP.Types (urlEncode)
import Test.Hspec (shouldSatisfy)
import System.Directory (removeDirectoryRecursive, doesDirectoryExist)
import qualified Data.ByteString.Char8 as BSC8
import qualified Text.Blaze.Html5 as H
import qualified Text.Blaze.Html5.Attributes as A
#endif

data SpecState where
  SpecState :: (ServerKeySet k) =>
    { ssRandomSource :: RandomSource
    , ssAuthSettings :: AuthCookieSettings
    , ssServerKeySet :: k
    , ssGenerateKey  :: IO ()
    } -> SpecState


main :: IO ()
main = do
  rs <- mkRandomSource drgNew 1000

  return SpecState
    { ssRandomSource = rs
    , ssAuthSettings = authSettings
    , ssServerKeySet = mkPersistentServerKey "0123456789abcdef"
    , ssGenerateKey  = return ()
    } >>= hspec . basicSpec

#if MIN_VERSION_servant (0,9,1) && MIN_VERSION_directory (1,2,5)
  let rmDir name = doesDirectoryExist name
        >>= \exists -> when exists $ removeDirectoryRecursive name

  bracket
    (do
       let keySetDir = "./test-key-set"
       rmDir keySetDir
       return FileKSParams
         { fkspMaxKeys = 3
         , fkspKeySize = 16
         , fkspPath = keySetDir
         } >>= \fksp -> (fksp,) <$> mkFileKeySet fksp)

    (rmDir . fkspPath . fst)

    (\(fksp, ks) -> hspec . renewalSpec $ SpecState
         { ssRandomSource = rs
         , ssAuthSettings = authSettings
         , ssServerKeySet = ks
         , ssGenerateKey  = mkFileKey fksp
         })
#endif

basicSpec :: SpecState -> Spec
basicSpec ss@(SpecState {..}) = describe "basic functionality" $ with
  (return $ app ssAuthSettings ssGenerateKey ssRandomSource ssServerKeySet) $ do

  let form = LoginForm {
          lfUsername = "mr_foo"
        , lfPassword = "password1"
        , lfRemember = False
        , lfRenew    = False
        }

  context "home page" $ do
    it "responds successfully" $ do
      get "/" `shouldRespondWithMarkup` homePage

  context "login page" $ do
    it "responds successfully" $ do
      get "/login" `shouldRespondWithMarkup` (loginPage True)

    it "shows message on incorrect login" $ do
      (login form { lfPassword = "wrong" }) `shouldRespondWithMarkup` (loginPage False)

    let hasExpirationFlags
          = not . null
          . filter ((`elem` ["Max-Age", "Expires"]) . fst)
          . parseCookies

    it "responds with session cookies if 'Remember me' is not set" $ do
      (login form { lfRemember = False }
        >>= return . hasExpirationFlags . getCookieValue)
        >>= liftIO . (`shouldBe` False)

    it "responds with normal cookies if `Remember me` is set" $ do
      (login form { lfRemember = True }
        >>= return . hasExpirationFlags . getCookieValue)
        >>= liftIO . (`shouldBe` True)

  context "private page" $ do
    it "rejects requests without cookies" $ do
      get "/private" `shouldRespondWith` 403 { matchBody = matchBody' "No cookies" }

    it "accepts requests with proper cookies" $ do
      (login form
         >>= return . getCookieValue
         >>= getPrivate) `shouldRespondWith` 200

    it "accepts requests with proper cookies (sanity check)" $ do
      (login form
        >>= liftIO . forgeCookies ss authSettings ssServerKeySet
        >>= getPrivate) `shouldRespondWith` 200

    it "rejects requests with incorrect MAC" $ do
      let newServerKeySet = mkPersistentServerKey "0000000000000000"
      (login form
        >>= liftIO . forgeCookies ss authSettings newServerKeySet
        >>= getPrivate) `shouldRespondWithException` (IncorrectMAC "")

    it "rejects requests with expired cookies" $ do
      let newAuthSettings = authSettings { acsMaxAge = 0 }
      let t = UTCTime (toEnum 0) 0
      (login form
        >>= liftIO . forgeCookies ss newAuthSettings ssServerKeySet
        >>= getPrivate) `shouldRespondWithException` (CookieExpired t t)

    let hasSetCookieHeader
          = maybe False (const True)
          . lookup "set-cookie"
          . simpleHeaders

    it "doesn't renew cookies when renew flag is not set" $ do
      (login (form { lfRemember = True, lfRenew = False })
        >>= return . getCookieValue
        >>= getPrivate
        >>= return . hasSetCookieHeader)
        >>= liftIO . (`shouldBe` False)

#if MIN_VERSION_servant(0,9,1)
    it "does renew cookies when renew flag is set" $ do
      (login form { lfRemember = True, lfRenew = True }
        >>= return . getCookieValue
        >>= getPrivate
        >>= return . hasSetCookieHeader)
        >>= liftIO . (`shouldBe` True)
#endif

#if MIN_VERSION_servant (0,9,1) && MIN_VERSION_directory (1,2,5)
renewalSpec :: SpecState -> Spec
renewalSpec (SpecState {..}) = describe "renewal functionality" $ with
  (return $ app ssAuthSettings ssGenerateKey ssRandomSource ssServerKeySet) $ do

  context "keys" $ do
    it "automatically creates a key" $ do
      keys <- extractKeys
      liftIO $ keys `shouldSatisfy` ((== 1) . length)

    it "adds new key" $ do
      keys <- extractKeys
      addKey
      keys' <- extractKeys
      liftIO $ keys `shouldBe` (tail keys')

    it "removes a key" $ do
      keys <- extractKeys
      remKey $ last keys
      keys' <- extractKeys
      liftIO $ keys' `shouldBe` (init keys)

  context "cookies" $ do
    let form = LoginForm {
            lfUsername = "mr_foo"
          , lfPassword = "password1"
          , lfRemember = False
          , lfRenew    = False
          }

    it "rejects requests with deleted keys" $ do
      cookieValue <- getCookieValue <$> login form
      getPrivate cookieValue `shouldRespondWith` 200

      key <- head <$> extractKeys
      addKey >> remKey key

      getPrivate cookieValue `shouldRespondWith` 403

    it "accepts requests with old key and renews cookie" $ do
      cookieValue <- getCookieValue <$> login form
      getPrivate cookieValue `shouldRespondWith` 200

      key <- head <$> extractKeys
      addKey
      newCookieValue <- getCookieValue <$> getPrivate cookieValue

      remKey key
      getPrivate newCookieValue `shouldRespondWith` 200

    it "does not renew cookies for the newest key" $ do
      cookieValue <- getCookieValue <$> login form
      _ <- getPrivate cookieValue `shouldRespondWith` 200
      r <- getPrivate cookieValue
      liftIO $ (lookup "set-cookie" $ simpleHeaders r) `shouldBe` Nothing
#endif

#if MIN_VERSION_hspec_wai (0,7,0)
matchBody' :: BSL.ByteString -> MatchBody
matchBody' = bodyEquals
#else
matchBody' :: BSL.ByteString -> Maybe BSL.ByteString
matchBody' = Just
#endif

#if MIN_VERSION_servant (0,9,0)
encode :: ToForm a => a -> BSL.ByteString
encode = urlEncodeForm . toForm
#else
encode :: ToFormUrlEncoded a => a -> BSL.ByteString
encode = mimeRender (Proxy :: Proxy FormUrlEncoded)
#endif

shrinkBody :: Int64 -> SResponse -> SResponse
shrinkBody len r = r { simpleBody = BSL.take len $ simpleBody r }

shouldRespondWithException :: WaiSession SResponse -> AuthCookieException -> WaiExpectation
shouldRespondWithException req ex = do
  let exception = BSLC8.pack . head . words . show $ ex
  (shrinkBody (BSLC8.length exception) <$> req) `shouldRespondWith` 403 {
      matchBody = matchBody' exception
    }

shouldRespondWithMarkup :: WaiSession SResponse -> Markup -> WaiExpectation
shouldRespondWithMarkup req markup = do
  req `shouldRespondWith` 200 {
      matchBody = matchBody' $ renderMarkup markup
    }

formContentType :: Header
formContentType = (
    hContentType
  , renderHeader $ contentType (Proxy :: Proxy FormUrlEncoded))

login :: LoginForm -> WaiSession SResponse
login lf = request
  methodPost "/login" [formContentType] (encode lf)

getPrivate :: BS.ByteString -> WaiSession SResponse
getPrivate cookieValue = request
  methodGet "/private" [(hCookie, cookieValue)] ""

extractSession :: SpecState -> SResponse -> IO (ExtendedPayloadWrapper Account)
extractSession SpecState {..} SResponse {..} = maybe
  (error "cookies aren't available")
  (decryptSession ssAuthSettings ssServerKeySet)
  (parseSessionResponse ssAuthSettings simpleHeaders)

forgeCookies :: (ServerKeySet k)
  => SpecState
  -> AuthCookieSettings
  -> k
  -> SResponse
  -> IO BS.ByteString
forgeCookies ss newAuthSettings newServerKeySet r = extractSession ss r
  >>= \s -> renderSession newAuthSettings (ssRandomSource ss) newServerKeySet def (epwSession s) ()

#if MIN_VERSION_servant (0,9,1) && MIN_VERSION_directory (1,2,5)
extractKeys :: WaiSession [BS.ByteString]
extractKeys = (extractKeys' . BSL.toStrict . simpleBody) <$> get "/keys" where
  del = '#'

  (openTag, closeTag) = (id *** BS.drop 1) $ BSC8.span (/= del) $
    BSL.toStrict . renderMarkup $
      H.span H.! A.class_ "key" $ H.toHtml [del]

  shrinkBy prefix = BS.drop . BS.length $ prefix

  extractKeys' body = let
    body' = snd $ BS.breakSubstring openTag body
    (key, rest) = shrinkBy openTag *** shrinkBy closeTag $
       BS.breakSubstring closeTag body'
    in if BS.null body'
       then []
       else key:(extractKeys' rest)

addKey :: WaiSession ()
addKey = void $ get "/keys/add"

remKey :: BS.ByteString -> WaiSession ()
remKey key = void $ get $ "/keys/rem/" <> (urlEncode True $ key)
#endif

getCookieValue :: SResponse -> BS.ByteString
getCookieValue = fromMaybe (error "cookies aren't available")
               . lookup "set-cookie"
               . simpleHeaders