content: {-# LANGUAGE DataKinds #-}
  {-# LANGUAGE DeriveAnyClass #-}
  {-# LANGUAGE DerivingStrategies #-}
  {-# LANGUAGE GeneralizedNewtypeDeriving #-}
  {-# LANGUAGE KindSignatures #-}
  {-# LANGUAGE RecordWildCards #-}
  {-# LANGUAGE TypeFamilies #-}
  {-# LANGUAGE TypeOperators #-}
  
  module Share.OAuth.Session
    ( Session (..),
      ServiceName (..),
      AdditionalSessionClaims (..),
      PendingSession (..),
      LoginRequest (..),
      AuthenticatedSession,
      AuthenticatedUserId,
      MaybeAuthenticatedSession,
      MaybeAuthenticatedUserId,
      pattern MaybeAuthedUserID,
      pattern AuthenticatedUser,
      pattern Unauthenticated,
      createSession,
      AuthCheckCtx,
      addAuthCheckCtx,
    )
  where
  
  import Control.Applicative
  import Control.Monad.Random
  import Control.Monad.Trans.Maybe (MaybeT (..))
  import Crypto.JWT qualified as JWT
  import Data.Aeson
  import Data.Aeson qualified as Aeson
  import Data.Binary
  import Data.Binary.Instances.Time ()
  import Data.ByteString qualified as BS
  import Data.Coerce (coerce)
  import Data.Kind (Type)
  import Data.Set (Set)
  import Data.Text (Text)
  import Data.Text qualified as Text
  import Data.Text.Encoding qualified as Text
  import Data.Time (NominalDiffTime, UTCTime, addUTCTime, getCurrentTime, nominalDay)
  import GHC.Generics (Generic)
  import Network.HTTP.Types qualified as Network
  import Network.URI
  import Network.Wai qualified as Wai
  import Servant
  import Servant.Client.Core.Auth qualified as ServantAuth
  import Servant.Server.Experimental.Auth qualified as ServantAuth
  import Share.JWT
  import Share.OAuth.Types
  import Share.Utils.API ((:++) (..))
  import Share.Utils.Binary
  import Share.Utils.IDs qualified as IDs
  import Share.Utils.Servant.Cookies qualified as Cookies
  import Web.Cookie (parseCookies)
  
  -- | Requires a valid session cookie to be present in the request and provides it as an
  -- argument to the handler
  --
  -- e.g. @Session -> ...@
  type AuthenticatedSession = Servant.AuthProtect "require-session"
  
  type instance ServantAuth.AuthServerData (Servant.AuthProtect "require-session") = Session
  
  type instance ServantAuth.AuthClientData (Servant.AuthProtect "require-session") = JWT.SignedJWT
  
  -- | Requires a valid session cookie to be present in the request,
  -- provides the authenticated user's user-id as an argument to the handler
  --
  -- e.g. @UserId -> ...@
  type AuthenticatedUserId = Servant.AuthProtect "require-user-id"
  
  type instance ServantAuth.AuthServerData (AuthProtect "require-user-id") = UserId
  
  type instance ServantAuth.AuthClientData (AuthProtect "require-user-id") = JWT.SignedJWT
  
  -- | Used for endpoints with optional auth.
  -- Provides 'Just' the session if a valid session cookie is present in the request,
  -- otherwise provides 'Nothing'.
  --
  -- e.g. @Maybe Session -> ...@
  type MaybeAuthenticatedSession = Servant.AuthProtect "maybe-session"
  
  type instance ServantAuth.AuthServerData (AuthProtect "maybe-session") = Maybe Session
  
  type instance ServantAuth.AuthClientData (AuthProtect "maybe-session") = Maybe JWT.SignedJWT
  
  -- | Used for endpoints with optional auth.
  -- Provides 'Just' the user ID if a valid session cookie is present in the request,
  -- otherwise provides 'Nothing'.
  type MaybeAuthenticatedUserId = Servant.AuthProtect "maybe-user-id"
  
  type instance ServantAuth.AuthServerData (AuthProtect "maybe-user-id") = Maybe UserId
  
  type instance ServantAuth.AuthClientData (AuthProtect "maybe-user-id") = Maybe JWT.SignedJWT
  
  -- | An additional check to perform on a session. Returns True if valid, False otherwise.
  type SessionCheck = (Session -> Handler Bool)
  
  -- | A custom auth handler which checks the session cookie and the JWT token and returns the
  -- authenticated session if there is one.
  -- If a valid session is authenticated then the additional session check is performed, if that
  -- check fails the user is considered unauthenticated.
  checkOptionalAuthenticatedSession ::
    Text ->
    SessionCheck ->
    Cookies.CookieSettings {- Not actually used yet, but will probably need it soon so this makes it future-compatible -} ->
    Share.JWT.JWTSettings ->
    ServantAuth.AuthHandler Wai.Request (Maybe Session)
  checkOptionalAuthenticatedSession sessionCookieName sessCheck _cookieSettings jwtSettings = ServantAuth.mkAuthHandler authHandler
    where
      authHandler :: Wai.Request -> Handler (Maybe Session)
      authHandler req = do
        let guardCheck :: Session -> MaybeT Handler Session
            guardCheck s = do
              lift (sessCheck s) >>= guard
              pure s
        runMaybeT $ (jwtSession req >>= guardCheck) <|> (cookieSession req >>= guardCheck)
      jwtSession :: Wai.Request -> MaybeT Handler Session
      jwtSession req = do
        tokenBytes <- MaybeT . pure $ do
          authHdr <- lookup "Authorization" $ Wai.requestHeaders req
          let bearer = "Bearer "
              (mbearer, rest) = BS.splitAt (BS.length bearer) authHdr
          guard (mbearer == bearer)
          pure rest
        signedJWT <- MaybeT . pure . eitherToMaybe $ Share.JWT.textToSignedJWT (Text.decodeUtf8 tokenBytes)
        MaybeT . liftIO . fmap eitherToMaybe $ Share.JWT.verifyJWT jwtSettings signedJWT
      cookieSession :: Wai.Request -> MaybeT Handler Session
      cookieSession req = do
        jwtCookieBytes <- MaybeT . pure $ do
          cookies' <- lookup Network.hCookie $ Wai.requestHeaders req
          let cookies = parseCookies cookies'
          lookup (Text.encodeUtf8 sessionCookieName) cookies
        signedJWT <- MaybeT . pure . eitherToMaybe $ Share.JWT.textToSignedJWT (Text.decodeUtf8 jwtCookieBytes)
        MaybeT . liftIO . fmap eitherToMaybe $ Share.JWT.verifyJWT jwtSettings signedJWT
  
      eitherToMaybe = \case
        Left _ -> Nothing
        Right a -> Just a
  
  -- | Make an auth handler using an additional session check function.
  checkRequiredAuthenticatedSession :: Text -> SessionCheck -> Cookies.CookieSettings -> Share.JWT.JWTSettings -> ServantAuth.AuthHandler Wai.Request (Session)
  checkRequiredAuthenticatedSession sessionCookieName sessCheck cookieSettings jwtSettings =
    ServantAuth.mkAuthHandler $ \req -> do
      ServantAuth.unAuthHandler (checkOptionalAuthenticatedSession sessionCookieName sessCheck cookieSettings jwtSettings) req >>= \case
        Nothing -> throwError (Servant.err401 {errReasonPhrase = "Unauthenticated", errBody = "Unauthenticated: Please ensure you have a valid session."})
        Just a -> pure a
  
  -- | The context args required for auth handlers.
  type AuthCheckCtx =
    '[ ServantAuth.AuthHandler Wai.Request Session,
       ServantAuth.AuthHandler Wai.Request UserId,
       ServantAuth.AuthHandler Wai.Request (Maybe Session),
       ServantAuth.AuthHandler Wai.Request (Maybe UserId)
     ]
  
  -- | Here we provide the handlers for each AuthRequirement type.
  --
  -- You can combine this with your app's context using '.++' at both the term and type
  -- level
  authCheckCtx ::
    Cookies.CookieSettings ->
    Share.JWT.JWTSettings ->
    Text ->
    Servant.Context AuthCheckCtx
  authCheckCtx cookieSettings jwtSettings sessionCookieName =
    (checkRequiredAuthenticatedSession sessionCookieName defaultSessionCheck cookieSettings jwtSettings)
      :. (mapAuthHandler sessionUserId $ checkRequiredAuthenticatedSession sessionCookieName defaultSessionCheck cookieSettings jwtSettings)
      :. (checkOptionalAuthenticatedSession sessionCookieName defaultSessionCheck cookieSettings jwtSettings)
      :. (mapAuthHandler (fmap sessionUserId) $ checkOptionalAuthenticatedSession sessionCookieName defaultSessionCheck cookieSettings jwtSettings)
      :. Servant.EmptyContext
    where
      -- By default we don't have any additional checks aside from those performed as part of
      -- JWT decoding.
      defaultSessionCheck _session = pure True
      -- A functor instance is added in 0.20, but we're on 0.19 for now.
      mapAuthHandler :: (a -> a') -> ServantAuth.AuthHandler req a -> ServantAuth.AuthHandler req a'
      mapAuthHandler f (ServantAuth.AuthHandler h) = ServantAuth.AuthHandler $ \req -> f <$> h req
  
  addAuthCheckCtx ::
    Cookies.CookieSettings ->
    Share.JWT.JWTSettings ->
    Text ->
    Servant.Context (appCtx :: [Type]) ->
    Servant.Context (AuthCheckCtx .++ appCtx)
  addAuthCheckCtx cookieSettings jwtSettings sessionCookieName appCtx =
    authCheckCtx cookieSettings jwtSettings sessionCookieName .++ appCtx
  
  -- | Unpacks an auth result into a Maybe, and extracts the user id.
  -- This is useful for pattern matching on the result of an auth check and extracting the user id.
  pattern MaybeAuthedUserID :: Maybe UserId -> Maybe Session
  pattern MaybeAuthedUserID mayUserId <- (getAuthenticatedUserId -> mayUserId)
  
  {-# COMPLETE MaybeAuthedUserID #-}
  
  -- | Only matches if a session was successfully authenticated
  pattern AuthenticatedUser :: UserId -> Maybe Session
  pattern AuthenticatedUser userId <- Just (Session {sessionUserId = userId})
  
  -- | Matches if the user is unauthenticated, regardless of the reason.
  pattern Unauthenticated :: Maybe Session
  pattern Unauthenticated <- MaybeAuthedUserID Nothing
  
  {-# COMPLETE AuthenticatedUser, Unauthenticated #-}
  
  getAuthenticatedUserId :: Maybe Session -> Maybe UserId
  getAuthenticatedUserId = \case
    Just ses -> Just $ sessionUserId ses
    _ -> Nothing
  
  data AdditionalSessionClaims = AdditionalSessionClaims
    { termsAndConditionsAccepted :: Bool
    }
    deriving stock (Show, Eq, Ord)
  
  -- | The service which proxied a user request to us, if any.
  -- E.g. A user may make a call to the cloud-api backend, which then may
  -- proxy that request to the share-api backend, in which case the ServiceName
  -- would be "cloud-api".
  newtype ServiceName = ServiceName Text
    deriving stock (Show, Eq, Ord)
    deriving newtype (FromJSON, ToJSON)
  
  -- | Note: This session will be signed into the user's cookie unencrypted.
  -- Don't put anything sensitive in here.
  data Session = Session
    { sessionId :: SessionId,
      sessionUserId :: UserId,
      sessionServiceName :: Maybe ServiceName,
      sessionCreated :: UTCTime,
      sessionExpiry :: UTCTime,
      sessionIssuer :: Issuer,
      sessionAudience :: Set Audience
    }
    deriving stock (Show)
    deriving (Binary) via JSONBinary Session
    deriving (AsJWTClaims) via JSONJWTClaims Session
  
  instance ToJSON Session where
    toJSON Session {..} =
      toJSON $
        StandardClaims
          { sub = IDs.toText sessionUserId,
            iat = sessionCreated,
            exp = sessionExpiry,
            iss = sessionIssuer,
            aud = sessionAudience,
            jti = IDs.toText (coerce @SessionId @JTI sessionId)
          }
          :++ Aeson.object ["serviceName" .= sessionServiceName]
  
  instance FromJSON Session where
    parseJSON v = do
      StandardClaims {..} :++ obj <- parseJSON v
      JTI sessionId <- either (fail . Text.unpack) pure $ IDs.fromText jti
      sessionUserId <- either (fail . Text.unpack) pure $ IDs.fromText sub
      sessionServiceName <- obj Aeson..:? "serviceName"
      pure Session {sessionId = SessionId sessionId, sessionUserId, sessionCreated = iat, sessionExpiry = exp, sessionIssuer = iss, sessionAudience = aud, sessionServiceName}
  
  data PendingSession = PendingSession
    { pendingId :: PendingSessionId,
      loginRequest :: LoginRequest,
      pkceVerifier :: PKCEVerifier,
      -- The URI that the user started the login flow from,
      -- you may wish to redirect them back here after login.
      returnToURI :: Maybe URI,
      -- Any additional data about the login/signup that the service provider wanted to pass
      -- through to the auth process.
      additionalData :: Maybe Value
    }
    deriving stock (Show, Generic)
    deriving anyclass (ToJSON, FromJSON)
    deriving (Binary) via JSONBinary PendingSession
  
  data LoginRequest
    = -- The login was kicked off by an oauth client, e.g. UCM
      OAuthRequest AuthenticationRequest
    | --  The login was kicked off by visiting the 'login' page.
      LoginPage
    deriving stock (Show, Generic)
    deriving anyclass (ToJSON, FromJSON)
    deriving (Binary) via JSONBinary LoginRequest
  
  -- | Sessions expire in 1 year
  sessionTTL :: NominalDiffTime
  sessionTTL =
    (365 * nominalDay)
  
  createSession :: (MonadIO m) => ServiceName -> Issuer -> Set Audience -> UserId -> m Session
  createSession serviceName sessionIssuer sessionAudience sessionUserId = do
    let sessionServiceName = Just serviceName
    sessionId <- randomIO
    sessionCreated <- liftIO getCurrentTime
    let sessionExpiry = addUTCTime sessionTTL sessionCreated
    let session = Session {..}
    pure session
  