module Nuthatch.Session
  ( mintSession
  , verifySession
  , newCsrfToken
  , constantTimeEq
  ) where

import           Control.Lens
  ( (&)
  , (?~)
  , (^.)
  , (.~)
  )
import           Crypto.Hash (Digest, SHA256, hash)
import           Crypto.JOSE (runJOSE)
import           Crypto.JOSE.Compact (decodeCompact, encodeCompact)
import           Crypto.JOSE.JWA.JWS (Alg (HS256))
import           Crypto.JOSE.JWK (fromOctets)
import           Crypto.JOSE.JWS
  ( algorithms
  , newJWSHeaderProtected
  )
import           Crypto.JWT
  ( Audience (..)
  , ClaimsSet
  , HasClaimsSet (..)
  , JWTError
  , NumericDate (..)
  , SignedJWT
  , defaultJWTValidationSettings
  , emptyClaimsSet
  , issuerPredicate
  , signJWT
  , verifyJWT
  )
import qualified Data.Aeson as A
import qualified Data.Aeson.KeyMap as KM
import qualified Data.ByteArray as BA
import qualified Data.ByteString as BS
import qualified Data.ByteString.Base64.URL as B64
import qualified Data.ByteString.Lazy as BL
import           Data.Set (singleton)
import           Data.String (fromString)
import           Data.Text (Text)
import qualified Data.Text as T
import qualified Data.Text.Encoding as TE
import           Data.Time.Clock (addUTCTime, getCurrentTime)
import           Nuthatch.Config (Config (..))
import           Nuthatch.Paths (encodeUserId)
import           Nuthatch.Types
import           System.Entropy (getEntropy)

data SessionClaims = SessionClaims
  { registeredClaims :: ClaimsSet
  , oidcIssuer       :: Text
  , oidcSubject      :: Text
  , displayName      :: Maybe Text
  }
  deriving stock (Eq, Show)

instance HasClaimsSet SessionClaims where
  claimsSet update claims =
    fmap
      (\registered -> claims { registeredClaims = registered })
      (update (registeredClaims claims))

instance A.FromJSON SessionClaims where
  parseJSON value = A.withObject "SessionClaims" parse value
    where
      parse object =
        SessionClaims
          <$> A.parseJSON value
          <*> object A..: "oidc_issuer"
          <*> object A..: "oidc_subject"
          <*> object A..:? "display"

instance A.ToJSON SessionClaims where
  toJSON claims =
    case A.toJSON (registeredClaims claims) of
      A.Object object ->
        A.Object $
          KM.insert "oidc_issuer" (A.String (oidcIssuer claims)) $
          KM.insert "oidc_subject" (A.String (oidcSubject claims)) $
          maybe
            object
            (\display -> KM.insert "display" (A.String display) object)
            (displayName claims)
      value -> value

mintSession :: Config -> Text -> Text -> Maybe Text -> IO (Either Text SessionToken)
mintSession cfg issuer subject display = do
  now <- getCurrentTime
  let expires = addUTCTime (fromInteger (cfgSessionHours cfg * 3600)) now
      userId = encodeUserId issuer subject
      applicationIssuer = "nuthatch"
      applicationAudience = "nuthatch"
      userSubject = fromString (T.unpack (unUserId userId))
  result <- runJOSE $ do
    let claims = SessionClaims
          { registeredClaims =
              emptyClaimsSet
                & claimIss ?~ applicationIssuer
                & claimSub ?~ userSubject
                & claimAud ?~ Audience [applicationAudience]
                & claimIat ?~ NumericDate now
                & claimExp ?~ NumericDate expires
          , oidcIssuer = issuer
          , oidcSubject = subject
          , displayName = display
          }
        key = fromOctets (TE.encodeUtf8 (cfgJwtSecret cfg))
        header = newJWSHeaderProtected HS256
    encodeCompact <$> signJWT key header claims
  pure $ case result of
    Left (_ :: JWTError) -> Left "session token creation failed"
    Right compact ->
      validateSessionToken (TE.decodeUtf8 (BL.toStrict compact))

verifySession :: Config -> SessionToken -> IO (Either Text SessionUser)
verifySession cfg (SessionToken token) = do
  let applicationIssuer = "nuthatch"
      applicationAudience = "nuthatch"
  result <- runJOSE $ do
    jwt <- decodeCompact (BL.fromStrict (TE.encodeUtf8 token))
    let settings =
          defaultJWTValidationSettings (== applicationAudience)
            & issuerPredicate .~ (== applicationIssuer)
            & algorithms .~ singleton HS256
        key = fromOctets (TE.encodeUtf8 (cfgJwtSecret cfg))
    verifyJWT settings key (jwt :: SignedJWT)
  pure $ case result of
    Left (_ :: JWTError) -> Left "session token validation failed"
    Right claims ->
      let userId = encodeUserId (oidcIssuer claims) (oidcSubject claims)
          expectedSubject = fromString (T.unpack (unUserId userId))
          registered = registeredClaims claims
      in if registered ^. claimSub /= Just expectedSubject
          || registered ^. claimAud /= Just (Audience ["nuthatch"])
          || registered ^. claimIat == Nothing
          || registered ^. claimExp == Nothing
        then Left "session token claims are invalid"
        else Right SessionUser
          { suIssuer = oidcIssuer claims
          , suSubject = oidcSubject claims
          , suUserId = userId
          , suDisplay = displayName claims
          }

newCsrfToken :: IO CsrfToken
newCsrfToken = do
  bytes <- getEntropy 32
  pure (CsrfToken (TE.decodeUtf8 (B64.encodeUnpadded bytes)))

constantTimeEq :: CsrfToken -> CsrfToken -> Bool
constantTimeEq (CsrfToken a) (CsrfToken b) =
  let left = TE.encodeUtf8 a
      right = TE.encodeUtf8 b
  in BS.length left == BS.length right
     && BA.constEq
          (hash left :: Digest SHA256)
          (hash right :: Digest SHA256)