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

import           Crypto.Hash (Digest, SHA256, hash)
import           Crypto.MAC.HMAC (HMAC, hmac)
import qualified Data.Aeson as A
import qualified Data.Aeson.Key as Key
import qualified Data.Aeson.KeyMap as KM
import           Data.ByteArray (convert)
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.Text (Text)
import qualified Data.Text as T
import qualified Data.Text.Encoding as TE
import           Data.Time.Clock.POSIX (getPOSIXTime)
import           Nuthatch.Config (Config (..))
import           Nuthatch.Paths (encodeUserId)
import           Nuthatch.Types
import           System.Entropy (getEntropy)

mintSession :: Config -> Text -> Text -> Maybe Text -> IO (Either Text SessionToken)
mintSession cfg iss sub display = do
  now <- floor <$> getPOSIXTime
  let expAt = now + cfgSessionHours cfg * 3600
      uid = encodeUserId iss sub
      header = A.object
        [ "alg" A..= ("HS256" :: Text)
        , "typ" A..= ("JWT" :: Text)
        ]
      payload = A.object $
        [ "iss" A..= ("nuthatch" :: Text)
        , "aud" A..= ("nuthatch" :: Text)
        , "iat" A..= now
        , "exp" A..= expAt
        , "sub" A..= unUserId uid
        , "oidc_issuer" A..= iss
        , "oidc_subject" A..= sub
        ] <> maybe [] (\d -> ["display" A..= d]) display
      token = signJwt (TE.encodeUtf8 (cfgJwtSecret cfg)) header payload
  pure (validateSessionToken token)

verifySession :: Config -> SessionToken -> IO (Either Text SessionUser)
verifySession cfg (SessionToken token) = do
  now <- floor <$> getPOSIXTime
  pure $ do
    claims <- verifyJwt (TE.encodeUtf8 (cfgJwtSecret cfg)) token
    expAt <- claimInt claims "exp"
    if expAt < now
      then Left "session expired"
      else do
        iss <- claimText claims "oidc_issuer"
        sub <- claimText claims "oidc_subject"
        let display = either (const Nothing) Just (claimText claims "display")
            uid = encodeUserId iss sub
        pure SessionUser
          { suIssuer = iss
          , suSubject = sub
          , suUserId = uid
          , suDisplay = display
          }

signJwt :: BS.ByteString -> A.Value -> A.Value -> Text
signJwt secret header payload =
  let h = b64Json header
      p = b64Json payload
      signingInput = h <> "." <> p
      sig = b64Raw (hmacSha256 secret signingInput)
  in TE.decodeUtf8 (signingInput <> "." <> sig)

verifyJwt :: BS.ByteString -> Text -> Either Text A.Object
verifyJwt secret token =
  case T.splitOn "." token of
    [h, p, s] -> do
      let signingInput = TE.encodeUtf8 (h <> "." <> p)
          expected = b64Raw (hmacSha256 secret signingInput)
      if not (BA.constEq expected (TE.encodeUtf8 s))
        then Left "invalid session signature"
        else do
          payloadBytes <- b64Decode (TE.encodeUtf8 p)
          case A.eitherDecodeStrict payloadBytes of
            Left err -> Left (T.pack err)
            Right (A.Object obj) -> Right obj
            Right _ -> Left "jwt payload must be an object"
    _ -> Left "malformed session token"

hmacSha256 :: BS.ByteString -> BS.ByteString -> BS.ByteString
hmacSha256 secret msg =
  convert (hmac secret msg :: HMAC SHA256)

b64Json :: A.Value -> BS.ByteString
b64Json = B64.encodeUnpadded . BL.toStrict . A.encode

b64Raw :: BS.ByteString -> BS.ByteString
b64Raw = B64.encodeUnpadded

b64Decode :: BS.ByteString -> Either Text BS.ByteString
b64Decode bs =
  case B64.decode bs of
    Left err -> Left (T.pack err)
    Right x  -> Right x

claimText :: A.Object -> Text -> Either Text Text
claimText obj key =
  case KM.lookup (Key.fromText key) obj of
    Just (A.String t) -> Right t
    _                 -> Left ("missing claim " <> key)

claimInt :: A.Object -> Text -> Either Text Integer
claimInt obj key =
  case KM.lookup (Key.fromText key) obj of
    Just v ->
      case A.fromJSON v of
        A.Success (i :: Integer) -> Right i
        _                        -> Left ("invalid integer claim " <> key)
    _ -> Left ("missing claim " <> key)

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 ba = TE.encodeUtf8 a
      bb = TE.encodeUtf8 b
  in BS.length ba == BS.length bb
     && BA.constEq
          (hash ba :: Digest SHA256)
          (hash bb :: Digest SHA256)