-- SPDX-FileCopyrightText: 2022 Gregor Kleen <gregor.kleen@ifi.lmu.de>,Sarah Vaupel <sarah.vaupel@ifi.lmu.de>,Sarah Vaupel <vaupel.sarah@campus.lmu.de>
--
-- SPDX-License-Identifier: AGPL-3.0-or-later

module Utils.Tokens
  ( bearerToken
  , encodeBearer, BearerTokenException(..), decodeBearer
  , bearerParseJSON'
  , askBearer
  , formEmbedBearerPost, formEmbedBearerGet
  , decodeUploadToken
  ) where

import Import.NoModel

import Utils.Lens

import Model
import Model.Tokens

import Jose.Jwk (JwkSet(..))
import qualified Jose.Jwt as Jose

import Data.Aeson.Types (Parser)
import qualified Data.Aeson as JSON
import qualified Data.Aeson.Parser as JSON
import qualified Data.Aeson.Parser.Internal as JSON (jsonEOF')
import qualified Data.Aeson.Internal as JSON (iparse, formatError)

import qualified Data.HashMap.Strict as HashMap

import Settings

import CryptoID

import Text.Blaze (Markup)


bearerParseJSON' :: forall site m.
                  ( Hashable (AuthId site), Eq (AuthId site)
                  , HasCryptoUUID (AuthId site) (ReaderT CryptoIDKey Parser)
                  , ParseRoute site 
                  , Hashable (Route site)
                  , MonadSite site m
                  , MonadCrypto m
                  , MonadCryptoKey m ~ CryptoIDKey
                  )
               => m (Value -> Parser (BearerToken site))
-- ^ Read `CryptoIDKey` from monadic context and construct a `Parser` for `BearerToken`s
bearerParseJSON' = do
  cidKey <- cryptoIDKey return
  return $ flip runReaderT cidKey . bearerParseJSON


bearerToken :: forall site m.
               ( MonadSite site m
               , MonadIO m
               , HasInstanceID site InstanceId
               , HasClusterID site ClusterId
               , HasAppSettings site
               )
            => HashSet (Either Value (AuthId site)) -- ^ Authority
            -> Maybe (AuthId site) -- ^ Impersonate
            -> HashMap BearerTokenRouteMode (HashSet (Route site))
            -> Maybe AuthDNF -- ^ Additional auth
            -> Maybe (Maybe UTCTime) -- ^ Expiration; @Nothing@ determines default expiry time automatically
            -> Maybe UTCTime -- ^ Start of Validity; @Nothing@ means token starts to be valid immediately
            -> m (BearerToken site)
-- ^ Smart constructor for `BearerToken`, does not set route restrictions (due to polymorphism), use `tokenRestrict`
bearerToken bearerAuthority bearerImpersonate bearerRoutes bearerAddAuth mBearerExpiresAt bearerStartsAt = do
  bearerIdentifier <- liftIO getRandom
  bearerIssuedAt <- liftIO getCurrentTime
  bearerIssuedBy <- getsSite $ view instanceID
  bearerIssuedFor <- getsSite $ view clusterID

  defaultExpiration <- getsSite $ view _appBearerExpiration
  
  let bearerExpiresAt
        | Just t <- mBearerExpiresAt
        = t
        | Just tDiff <- defaultExpiration
        = Just $ tDiff `addUTCTime` fromMaybe bearerIssuedAt bearerStartsAt
        | otherwise
        = Nothing
      bearerRestrictions = HashMap.empty

  return BearerToken{..}


encodeBearer :: forall site m.
               ( MonadSite site m
               , MonadIO m
               , HasJSONWebKeySet site JwkSet
               , HasAppSettings site
               , HasCryptoUUID (AuthId site) m
               , RenderRoute site
               )
            => BearerToken site -> m Jwt
-- ^ Call `bearerToJSON` and encode the result as a `Jwt` according to `appJwtEncoding`
encodeBearer token = do
  payload <- Jose.Claims . toStrict . JSON.encode <$> bearerToJSON token
  JwkSet jwks <- getsSite $ view jsonWebKeySet
  jwtEncoding <- getsSite $ view _appBearerEncoding
  throwLeft =<< liftIO (Jose.encode jwks jwtEncoding payload)


data BearerTokenException
  = BearerTokenJwtError Jose.JwtError -- ^ An Error occurred in the underlying `Jwt`-Implementation
  | BearerTokenUnsecured -- ^ `Jwt` is insufficiently secured (unsigned and not encrypted)
  | BearerTokenInvalidFormat String -- ^ Content of the `Jwt` could not be parsed as a `BearerToken`
  | BearerTokenExpired | BearerTokenNotStarted | BearerTokenWrongAudience
  deriving (Eq, Show, Generic)
  deriving anyclass (Exception)

decodeBearer :: forall site m.
               ( MonadSite site m
               , MonadIO m
               , HasJSONWebKeySet site JwkSet
               , Hashable (AuthId site), Eq (AuthId site)
               , HasCryptoUUID (AuthId site) (ReaderT CryptoIDKey Parser)
               , MonadCryptoKey m ~ CryptoIDKey
               , MonadCrypto m
               , ParseRoute site
               , Hashable (Route site)
               , HasAppSettings site, HasClusterID site ClusterId
               )
            => Jwt -> m (BearerToken site)
-- ^ Decode a `Jwt` and call `bearerParseJSON`
--
-- Throws `BearerTokenException`s
decodeBearer (Jwt bs) = do
  JwkSet jwks <- getsSite $ view jsonWebKeySet
  content <- either (throwM . BearerTokenJwtError) return =<< liftIO (Jose.decode jwks Nothing bs)
  content' <- case content of
    Jose.Unsecured _ -> throwM BearerTokenUnsecured
    Jose.Jws (_header, payload) -> return payload
    Jose.Jwe (_header, payload) -> return payload
  parser <- bearerParseJSON'
  bearer@BearerToken{..} <- either (throwM . BearerTokenInvalidFormat . uncurry JSON.formatError) return $ JSON.eitherDecodeStrictWith JSON.jsonEOF' (JSON.iparse parser) content'
  bearerIssuedFor' <- getsSite $ view clusterID
  unless (bearerIssuedFor' == bearerIssuedFor) $
    throwM BearerTokenWrongAudience
  now <- liftIO getCurrentTime
  (clockLeniencyStart, clockLeniencyEnd) <- getsSite $ (,) <$> view _appBearerTokenClockLeniencyStart <*> view _appBearerTokenClockLeniencyEnd
  unless (NTop bearerExpiresAt > NTop (Just $ maybe id addUTCTime (negate <$> clockLeniencyEnd) now)) $
    throwM BearerTokenExpired
  unless (bearerStartsAt <= Just (maybe id addUTCTime clockLeniencyStart now)) $
    throwM BearerTokenNotStarted
  return bearer


askBearer :: forall m. MonadHandler m
          => m (Maybe Jwt)
-- ^ Retrieve current `Jwt` from HTTP-Header, POST-Parameter, or GET-Parameter
askBearer = runMaybeT $ asum
  [ MaybeT lookupBearerAuth >>= hoistMaybe . fromPathPiece
  , MaybeT $ lookupGlobalPostParam PostBearer
  , MaybeT $ lookupGlobalGetParam GetBearer
  , fmap Jwt . MaybeT $ lookupSessionBS (toPathPiece SessionBearer)
  ]

formEmbedBearerPost, formEmbedBearerGet :: MonadHandler m => (Markup -> m a) -> (Markup -> m a)
formEmbedBearerPost f fragment = do
  mBearer <- askBearer
  f [shamlet|
       $newline never
       $maybe bearer <- mBearer
         <input type=hidden name=#{toPathPiece PostBearer} value=#{toPathPiece bearer}>
       #{fragment}
    |]
formEmbedBearerGet f fragment = do
  mBearer <- askBearer
  f [shamlet|
       $newline never
       $maybe bearer <- mBearer
         <input type=hidden name=#{toPathPiece GetBearer} value=#{toPathPiece bearer}>
       #{fragment}
    |]


data UploadTokenException
  = UploadTokenJwtError Jose.JwtError -- ^ An Error occurred in the underlying `Jwt`-Implementation
  | UploadTokenUnsecured -- ^ `Jwt` is insufficiently secured (unsigned and not encrypted)
  | UploadTokenInvalidFormat String -- ^ Content of the `Jwt` could not be parsed as a `UploadToken`
  | UploadTokenExpired | UploadTokenNotStarted | UploadTokenWrongAudience
  deriving (Eq, Show, Generic)
  deriving anyclass (Exception)

decodeUploadToken :: forall m.
                     ( MonadHandler m, MonadThrow m
                     , HasJSONWebKeySet (HandlerSite m) JwkSet
                     , HasAppSettings (HandlerSite m), HasClusterID (HandlerSite m) ClusterId
                     )
                  => Jwt -> m UploadToken
-- ^ Decode a `Jwt`
--
-- Throws `UploadTokenException`s
decodeUploadToken (Jwt bs) = do
  JwkSet jwks <- getsYesod $ view jsonWebKeySet
  content <- either (throwM . UploadTokenJwtError) return =<< liftIO (Jose.decode jwks Nothing bs)
  content' <- case content of
    Jose.Unsecured _ -> throwM UploadTokenUnsecured
    Jose.Jws (_header, payload) -> return payload
    Jose.Jwe (_header, payload) -> return payload

  uploadToken@UploadToken{..} <- either (throwM . UploadTokenInvalidFormat) return $ JSON.eitherDecodeStrict content'

  uploadTokenIssuedFor' <- getsYesod $ view clusterID
  unless (uploadTokenIssuedFor' == uploadTokenIssuedFor) $
    throwM UploadTokenWrongAudience
  now <- liftIO getCurrentTime
  (clockLeniencyStart, clockLeniencyEnd) <- getsYesod $ (,) <$> view _appUploadTokenClockLeniencyStart <*> view _appUploadTokenClockLeniencyEnd
  unless (NTop uploadTokenExpiresAt > NTop (Just $ maybe id addUTCTime (negate <$> clockLeniencyEnd) now)) $
    throwM UploadTokenExpired
  unless (uploadTokenStartsAt <= Just (maybe id addUTCTime clockLeniencyStart now)) $
    throwM UploadTokenNotStarted
  return uploadToken