{-# LANGUAGE OverloadedStrings #-}

-- This file is part of the Wire Server implementation.
--
-- Copyright (C) 2025 Wire Swiss GmbH <opensource@wire.com>
--
-- This program is free software: you can redistribute it and/or modify it under
-- the terms of the GNU Affero General Public License as published by the Free
-- Software Foundation, either version 3 of the License, or (at your option) any
-- later version.
--
-- This program is distributed in the hope that it will be useful, but WITHOUT
-- ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
-- FOR A PARTICULAR PURPOSE. See the GNU Affero General Public License for more
-- details.
--
-- You should have received a copy of the GNU Affero General Public License along
-- with this program. If not, see <https://www.gnu.org/licenses/>.

module SAML2.WebSSO.SP where

import Control.Lens hiding (Level)
import Control.Monad
import Control.Monad.Except
import Control.Monad.Extra (ifM)
import Control.Monad.IO.Class
import Control.Monad.Reader
import Control.Monad.Writer
import Data.Foldable (toList)
import Data.Kind (Type)
import Data.List (nub, partition)
import Data.List.NonEmpty (NonEmpty)
import Data.Maybe (fromMaybe, isJust)
import Data.String.Conversions
import Data.Time
import Data.UUID (UUID)
import Data.UUID qualified as UUID
import Data.UUID.V4 qualified as UUID
import GHC.Stack
import SAML2.Util
import SAML2.WebSSO.API.UnvalidatedSAMLStatus
import SAML2.WebSSO.Config
import SAML2.WebSSO.Types
import Servant hiding (URI (..))
import System.Logger (Level (..))
import URI.ByteString

----------------------------------------------------------------------
-- class

-- | Application logic of the service provider.
type SP m = (HasConfig m, HasLogger m, HasCreateUUID m, HasNow m)

class HasLogger m where
  logger :: Level -> String -> m ()
  default logger :: (HasConfig m, MonadIO m) => Level -> String -> m ()
  logger = loggerConfIO

class HasCreateUUID m where
  createUUID :: m UUID
  default createUUID :: (MonadIO m) => m UUID
  createUUID = createUUIDIO

class HasNow m where
  getNow :: m Time
  default getNow :: (MonadIO m) => m Time
  getNow = getNowIO

type SPStore m = (SP m, SPStoreRequest AuthnRequest m, SPStoreAssertion Assertion m)

class SPStoreAssertion i m where
  storeAssertionInternal :: ID i -> Time -> m ()
  unStoreAssertion :: ID i -> m ()
  isAliveAssertion ::
    ID i ->
    -- | stored and not timed out.
    m Bool

class SPStoreRequest i m where
  storeRequest :: ID i -> Issuer {- NB: idp! -} -> Time -> m ()
  unStoreRequest :: ID i -> m ()
  getIdpIssuer :: ID i -> m (Maybe Issuer)

class (MonadError err m, Show (IdPConfigExtra m)) => SPStoreIdP err m where
  type IdPConfigExtra m :: Type
  type IdPConfigSPId m :: Type
  storeIdPConfig :: IdPConfig (IdPConfigExtra m) -> m ()
  getIdPConfig :: IdPId -> m (IdPConfig (IdPConfigExtra m))
  getIdPConfigByIssuer :: Issuer -> IdPConfigSPId m -> m (IdPConfig (IdPConfigExtra m))
  getIdPConfigByIssuer issuer spid = getIdPConfigByIssuerOptionalSPId issuer (Just spid)
  getIdPConfigByIssuerOptionalSPId :: Issuer -> Maybe (IdPConfigSPId m) -> m (IdPConfig (IdPConfigExtra m))

-- | HTTP handling of the service provider.
class (SP m, SPStore m, SPStoreIdP err m, MonadError err m) => SPHandler err m where
  type NTCTX m :: Type
  nt :: forall x. NTCTX m -> m x -> Handler x

----------------------------------------------------------------------
-- combinators

-- | Store 'Assertion's to prevent replay attack.  'Time' argument is end of life (IDs may be
-- garbage collected after that time).  Iff assertion has already been stored and is still alive,
-- return 'False'.
storeAssertion :: (Monad m, SPStore m) => ID Assertion -> Time -> m Bool
storeAssertion item endOfLife =
  ifM
    (isAliveAssertion item)
    (pure False)
    (True <$ storeAssertionInternal item endOfLife)

loggerConfIO :: (HasConfig m, MonadIO m) => Level -> String -> m ()
loggerConfIO level msg = do
  cfgsays <- (^. cfgLogLevel) <$> getConfig
  loggerIO cfgsays level msg

loggerIO :: (MonadIO m) => Level -> Level -> String -> m ()
loggerIO cfgsays level msg =
  when (level >= cfgsays) $ liftIO $ putStrLn msg

createUUIDIO :: (MonadIO m) => m UUID
createUUIDIO = liftIO UUID.nextRandom

getNowIO :: (MonadIO m) => m Time
getNowIO = Time <$> liftIO getCurrentTime

-- | (Microsoft Active Directory likes IDs to be of the form @id<32 hex digits>@: @ID . cs . ("id"
-- <>) . filter (/= '-') . cs . UUID.toText <$> createUUID@.  Hopefully the more common form
-- produced by this function is also ok.)
createID :: (Functor m, SP m) => m (ID a)
createID = ID . ("_" <>) . UUID.toText <$> createUUID

-- | Generate an 'AuthnRequest' value for the initiate-login response.  The 'NameIdPolicy' is
-- 'NameIDFUnspecified'.
--
-- We do not use XML encryption (which appears to have been dead for years).  We do not sign
-- 'AuthnRequest' values.  Both encryption and signatures are optional, and rarely found in the
-- wild.  Security: (1) the 'AuthnReq' 'ID' is stored, and 'Assertion's need to refer to a valid one
-- in order to get a positive 'AccessVerdict' by 'judge'.  'AuthnResponse's answering requests not
-- originating from us will be ignored.  (2) the request 'Issuer' is there to help the IdP construct
-- an 'AuthnResponse' that we will accept.  If it is changed by an attacker (either in the browser
-- or on the wire, which is weakly procted by TLS) and passed on to the legitimate IdP, there will
-- either be no access-granting 'AuthnResponse', or it will contain the wrong audience and be
-- rejected by us.  (3) The nameID policy is expected to be configured on the IdP side to not
-- support any set of name spaces that overlap (e.g. because user A has an email that is the account
-- name of user B).
createAuthnRequest :: (Monad m, SP m, SPStore m) => NominalDiffTime -> Issuer {- sp -} -> Issuer {- idp -} -> m AuthnRequest
createAuthnRequest lifeExpectancySecs _rqIssuer idpIssuer = do
  _rqID <- createID
  _rqIssueInstant <- getNow
  let _rqNameIDPolicy = Just $ NameIdPolicy NameIDFUnspecified Nothing True
  storeRequest _rqID idpIssuer (addTime lifeExpectancySecs _rqIssueInstant)
  pure AuthnRequest {..}

-- | The clock drift between IdP and us that we allow for.
--
-- FUTUREWORK: make this configurable
tolerance :: NominalDiffTime
tolerance = 60 -- seconds; must be positive

-- | First arg must be comfortably earlier than the latter: if 'tolerance' is one minute, then
-- @("4:32:14" `earlier` "4:31:15") == True@.
--
-- If this returns 'False' you should be *more* inclined to throw an error, so a good calling
-- pattern is @unless (a `earlier` b) throwSomething@, which is very different from @when (b
-- `earlier` a) throwSomething@.
earlier :: Time -> Time -> Bool
earlier early late = early < addTime tolerance late

-- | Even though this makes little sense, the standard requires to distinguish between "less"
-- and "less or equal".  For all practical purposes, you may consider @noLater == earlier@.
noLater :: Time -> Time -> Bool
noLater early late = early <= addTime tolerance late

----------------------------------------------------------------------
-- paths

-- | This function exists to deal with legacy test cases.
getSsoURINoMultiIngress ::
  forall m endpoint api.
  ( HasCallStack,
    Functor m,
    HasConfig m,
    IsElem endpoint api,
    HasLink endpoint,
    ToHttpApiData (MkLink endpoint Link)
  ) =>
  Proxy api ->
  Proxy endpoint ->
  m URI
getSsoURINoMultiIngress proxyAPI proxyAPIAuthResp =
  (extpath . _cfgSPSsoURI) <$> getMultiIngressDomainConfigNoMultiIngress
  where
    extpath :: URI -> URI
    extpath = (=/ (cs . toUrlPiece $ safeLink proxyAPI proxyAPIAuthResp))

-- | DANGER: This function is not valid for all spar configurations! It
-- spuriously fails for multi-ingress configs!
getMultiIngressDomainConfigNoMultiIngress :: forall m. (HasConfig m, Functor m) => m MultiIngressDomainConfig
-- FUTUREWORK: Get rid of this dangerous function. It solely exists to
-- implement legacy service example functions.
getMultiIngressDomainConfigNoMultiIngress =
  (fromMaybe (error "Configuration not found. (Multi-ingress config not supported.)") . (`getMultiIngressDomainConfig` Nothing))
    <$> getConfig

----------------------------------------------------------------------
-- compute access verdict(s)

-- | This monad collects errors in a writer, so that the reasons for access denial are as helpful as
-- possible.  It is a little like an exception monad, except you can throw several exceptions
-- without interrupting the flow, and will get a list of them at the end.
--
-- NOTE: @-XGeneralizedNewtypeDeriving@ does not help with the boilerplate instances below, since
-- this is a transformer stack and not a concrete 'Monad'.
newtype JudgeT m a = JudgeT
  {fromJudgeT :: ExceptT DeniedReason (WriterT [DeniedReason] (ReaderT JudgeCtx m)) a}

-- | Note on security: we assume that the SP has only one audience, which is defined here.  If you
-- have different sub-services running on your SP, associate a dedicated IdP with each sub-service.
-- (To be more specific, construct 'AuthnReq's to different IdPs for each sub-service.)  Secure
-- association with the service can then be guaranteed via the 'Issuer' in the signed 'Assertion'.
data JudgeCtx = JudgeCtx
  { _judgeCtxAudience :: Issuer,
    _judgeCtxResponseURI :: URI
  }
  deriving (Eq, Show)

makeLenses ''JudgeCtx

runJudgeT :: forall m. (Monad m, SP m) => JudgeCtx -> JudgeT m AccessVerdict -> m AccessVerdict
runJudgeT ctx (JudgeT em) = fmap collectErrors . (`runReaderT` ctx) . runWriterT $ runExceptT em
  where
    collectErrors :: (Either DeniedReason AccessVerdict, [DeniedReason]) -> AccessVerdict
    collectErrors (Left err, errs') = AccessDenied $ err : errs'
    collectErrors (Right _, errs@(_ : _)) = AccessDenied errs
    collectErrors (Right v, []) = v

-- the parts of the MonadError, MonadWriter interfaces we want here.
class (Functor m, Applicative m, Monad m) => MonadJudge m where
  getJudgeCtx :: m JudgeCtx
  deny :: DeniedReason -> m ()
  giveup :: DeniedReason -> m a

instance (Functor m, Applicative m, Monad m) => MonadJudge (JudgeT m) where
  getJudgeCtx = JudgeT . lift . lift $ ask
  deny = JudgeT . tell . (: [])
  giveup = JudgeT . throwError

instance (Functor m, Applicative m, Monad m) => Functor (JudgeT m) where
  fmap f = JudgeT . fmap f . fromJudgeT

instance (Functor m, Applicative m, Monad m) => Applicative (JudgeT m) where
  pure = JudgeT . pure
  (JudgeT f) <*> (JudgeT x) = JudgeT (f <*> x)

instance (Functor m, Applicative m, Monad m) => Monad (JudgeT m) where
  (JudgeT x) >>= f = JudgeT (x >>= fromJudgeT . f)

instance (Monad m, HasConfig m) => HasConfig (JudgeT m) where
  getConfig = JudgeT . lift . lift . lift $ getConfig

instance (Monad m, HasLogger m) => HasLogger (JudgeT m) where
  logger level = JudgeT . lift . lift . lift . logger level

instance (Monad m, HasCreateUUID m) => HasCreateUUID (JudgeT m) where
  createUUID = JudgeT . lift . lift . lift $ createUUID

instance (Monad m, HasNow m) => HasNow (JudgeT m) where
  getNow = JudgeT . lift . lift . lift $ getNow

instance (Monad m, SPStoreAssertion i m) => SPStoreAssertion i (JudgeT m) where
  storeAssertionInternal item = JudgeT . lift . lift . lift . storeAssertionInternal item
  unStoreAssertion = JudgeT . lift . lift . lift . unStoreAssertion
  isAliveAssertion = JudgeT . lift . lift . lift . isAliveAssertion

instance (Monad m, SPStoreRequest i m) => SPStoreRequest i (JudgeT m) where
  storeRequest item issuer = JudgeT . lift . lift . lift . storeRequest item issuer
  unStoreRequest = JudgeT . lift . lift . lift . unStoreRequest
  getIdpIssuer = JudgeT . lift . lift . lift . getIdpIssuer

-- | [3/4.1.4.2], [3/4.1.4.3]; specific to active-directory:
-- <https://docs.microsoft.com/en-us/azure/active-directory/develop/active-directory-single-sign-on-protocol-reference>
--
-- 'judge' does not consider the following parts of the 'AuthnResponse'.
--  * 'subjectID'
--  * 'scdAddress' ("If any bearer <SubjectConfirmationData> includes an Address attribute, the
--    service provider MAY check the user agent's client address against it."  [3/4.1.4.3])
--  * 'astSessionIndex'
--  * 'astSubjectLocality' ("This element is entirely advisory, since both of these fields are
--    quite easily “spoofed,” but may be useful information in some applications." [1/2.7.2.1])
--
-- 'judge' does *not* check any of the *unsigned* parts of the AuthnResponse body because...
-- those are not signed!  the standard doesn't seem to worry about that, but wire does.  this
-- affects `rspStatus`, `inRespTo`, `rspIssueInstant`, `rspDestination`.  those are inferred
-- from the *signed* information available.
judge :: (Monad m, SP m, SPStore m) => NonEmpty Assertion -> UnvalidatedSAMLStatus -> JudgeCtx -> m AccessVerdict
judge assertions status ctx = runJudgeT ctx $ do
  unless
    (eqUnvalidatedSAMLStatus status StatusSuccess)
    (deny DeniedStatusFailure)
  foldJudge assertions

foldJudge :: (HasCallStack, MonadJudge m, SP m, SPStore m) => NonEmpty Assertion -> m AccessVerdict
foldJudge (toList -> assertions) = do
  verdicts <- judge1 `mapM` assertions
  let (granteds, denieds) =
        verdicts
          & partition
            ( \case
                AccessDenied _ -> False
                AccessGranted _ -> True
            )

  pure $ case (granteds, denieds) of
    -- granted
    (nub -> [result@(AccessGranted _)], []) -> result
    -- denied
    ([], _ : _) -> AccessDenied (view avReasons =<< denieds)
    -- weird corner cases
    (bad@(_ : _), _) -> AccessDenied (DeniedBadUserRefs (show bad) : (view avReasons =<< denieds))
    ([], []) -> AccessDenied [DeniedBadUserRefs "there are no assertions"]

judge1 :: (HasCallStack, MonadJudge m, SP m, SPStore m) => Assertion -> m AccessVerdict
judge1 assertion = do
  inRespTo <- either (const $ giveup DeniedNoInResponseTo) pure $ assertionToInResponseTo assertion
  checkInResponseTo "response" (assertion ^. assIssuer) inRespTo
  verdict <- checkAssertion assertion
  unStoreRequest inRespTo
  pure verdict

-- | If this fails, we could continue ('deny'), but we stop processing ('giveup') to make DOS
-- attacks harder.
checkInResponseTo :: (SPStore m, MonadJudge m) => String -> Issuer -> ID AuthnRequest -> m ()
checkInResponseTo loc issuerFromRes req = do
  mbIssuerFromReq <- getIdpIssuer req
  unless
    (isJust mbIssuerFromReq)
    (giveup . DeniedBadInResponseTos $ loc <> ": " <> show req)
  unless
    (Just issuerFromRes == mbIssuerFromReq)
    (giveup (DeniedIssuerMismatch mbIssuerFromReq issuerFromRes))

checkIsInPast :: (SP m, MonadJudge m) => (Time -> Time -> DeniedReason) -> Time -> m ()
checkIsInPast err tim = do
  now <- getNow
  unless (tim `earlier` now) . deny $ err tim now

-- | Check that the response is intended for us (based on config's finalize-login uri stored in
-- 'JudgeCtx').
checkDestination :: (HasConfig m, MonadJudge m) => (String -> String -> DeniedReason) -> URI -> m ()
checkDestination err (renderURI -> expectedByIdp) = do
  (renderURI -> expectedByUs) <- (^. judgeCtxResponseURI) <$> getJudgeCtx
  unless (expectedByUs == expectedByIdp) . deny $ err (cs expectedByUs) (cs expectedByIdp)

checkAssertion :: (SP m, SPStore m, MonadJudge m) => Assertion -> m AccessVerdict
checkAssertion ass = do
  checkIsInPast DeniedAssertionIssueInstantNotInPast (ass ^. assIssueInstant)
  storeAssertion (ass ^. assID) (ass ^. assEndOfLife) >>= \case
    True -> pure ()
    False -> deny DeniedStatusFailure
  checkConditions `mapM_` (ass ^. assConditions)
  checkSubjectConfirmations ass
  let statements = ass ^. assContents . sasStatements
  unless (any isAuthnStatement statements) (deny DeniedNoAuthnStatement)
  checkStatement `mapM_` toList statements
  pure $ AccessGranted (assertionToUserRef ass)

checkStatement :: (SP m, MonadJudge m) => Statement -> m ()
checkStatement stm =
  do
    let issued = stm ^. astAuthnInstant
        mtimeout = stm ^. astSessionNotOnOrAfter
    checkIsInPast DeniedAuthnStatementIssueInstantNotInPast issued
    forM_ mtimeout $ \endoflife -> do
      now <- getNow
      unless (now `earlier` endoflife) . deny $ DeniedAuthnStatmentExpiredAt endoflife

-- | Check all 'SubjectConfirmation's and 'Subject's in all 'Assertion'.  Deny if not at least one
-- confirmation has method "bearer".
checkSubjectConfirmations :: (SP m, SPStore m, MonadJudge m) => Assertion -> m ()
checkSubjectConfirmations assertion = do
  bearerFlags :: [HasBearerConfirmation] <- case assertion ^. assContents . sasSubject of
    Subject _ confs -> checkSubjectConfirmation assertion `mapM` confs
  unless (nub bearerFlags == [HasBearerConfirmation]) $
    deny DeniedNoBearerConfSubj

data HasBearerConfirmation = HasBearerConfirmation | NoBearerConfirmation
  deriving (Eq, Ord, Bounded, Enum)

-- | Locally check one 'SubjectConfirmation' and deny if there is a problem.  If this is a
-- confirmation of method "bearer", return 'HasBearerConfirmation'.
checkSubjectConfirmation :: (SPStore m, MonadJudge m) => Assertion -> SubjectConfirmation -> m HasBearerConfirmation
checkSubjectConfirmation ass conf = do
  let bearer =
        if (conf ^. scMethod) == SubjectConfirmationMethodBearer
          then HasBearerConfirmation
          else NoBearerConfirmation
  when (bearer == HasBearerConfirmation) $ do
    when (all (null . (^. condAudienceRestriction)) (ass ^. assConditions)) $
      deny DeniedBearerConfAssertionsWithoutAudienceRestriction
  -- (the actual validation of the audience restrictions happens in 'checkConditions'.)

  checkSubjectConfirmationData `mapM_` (conf ^. scData)
  pure bearer

checkSubjectConfirmationData ::
  (HasConfig m, SP m, SPStore m, MonadJudge m) =>
  SubjectConfirmationData ->
  m ()
checkSubjectConfirmationData confdat = do
  checkDestination DeniedBadRecipient $ confdat ^. scdRecipient
  now <- getNow
  unless (now `earlier` (confdat ^. scdNotOnOrAfter)) . deny $
    DeniedNotOnOrAfterSubjectConfirmation (confdat ^. scdNotOnOrAfter)
  case confdat ^. scdNotBefore of
    Just notbef -> unless (notbef `noLater` now) . deny $ DeniedNotBeforeSubjectConfirmation notbef
    _ -> pure ()

checkConditions :: forall m. (HasCallStack, MonadJudge m, SP m) => Conditions -> m ()
checkConditions (Conditions lowlimit uplimit _onetimeuse audiences) = do
  now <- getNow
  case lowlimit of
    Just lim -> unless (lim `noLater` now) . deny $ DeniedNotBeforeCondition lim
    _ -> pure ()
  case uplimit of
    Just lim -> unless (now `earlier` lim) . deny $ DeniedNotOnOrAfterCondition lim
    _ -> pure ()
  Issuer us <- (^. judgeCtxAudience) <$> getJudgeCtx
  let checkAudience :: NonEmpty URI -> m ()
      checkAudience aus =
        unless (us `elem` aus) . deny $
          DeniedAudienceMismatch us aus
  checkAudience `mapM_` audiences
