module Nuthatch.Probe.Oidc
  ( ProbeContext (..)
  , Session (..)
  , login
  , newProbeContext
  , verifyPkceBoundary
  ) where

import           Control.Concurrent (threadDelay)
import           Control.Exception (SomeException, try)
import           Data.Aeson
  ( FromJSON (..)
  , eitherDecodeStrict'
  , withObject
  , (.:)
  )
import qualified Data.ByteString as BS
import qualified Data.ByteString.Char8 as B8
import           Data.Char (toLower)
import           Data.List (find)
import           Data.Maybe (fromMaybe)
import           Nuthatch.Probe.Assert
import           Nuthatch.Probe.Http
import           System.Environment (lookupEnv)

data ProbeContext = ProbeContext
  { probeClient           :: ProbeClient
  , probeBase             :: String
  , probeOidcIssuer       :: String
  , probeOidcProvider     :: String
  , probeOidcPassword     :: String
  , probeOidcLoginPat     :: Maybe String
  , probeOidcClientId     :: String
  , probeOidcClientSecret :: String
  , probeOidcRedirectUri  :: String
  }

data Session = Session
  { sessionCookies :: CookieSession
  , sessionCsrf    :: BS.ByteString
  }

data ZitadelSession = ZitadelSession
  { zitadelSessionId    :: String
  , zitadelSessionToken :: String
  }

instance FromJSON ZitadelSession where
  parseJSON = withObject "ZitadelSession" $ \object ->
    ZitadelSession
      <$> object .: "sessionId"
      <*> object .: "sessionToken"

newtype ZitadelCallback = ZitadelCallback String

instance FromJSON ZitadelCallback where
  parseJSON = withObject "ZitadelCallback" $ \object ->
    ZitadelCallback <$> object .: "callbackUrl"

newtype OidcDiscovery = OidcDiscovery String

instance FromJSON OidcDiscovery where
  parseJSON = withObject "OidcDiscovery" $ \object ->
    OidcDiscovery <$> object .: "token_endpoint"

newProbeContext :: IO ProbeContext
newProbeContext = do
  client <- newProbeClient
  lane <- requireEnv "NUTHATCH_FIXTURE_LANE"
  issuer <- requireEnv "NUTHATCH_OIDC_ISSUER"
  provider <- requireEnv "NUTHATCH_OIDC_PROVIDER"
  password <- requireEnv "NUTHATCH_OIDC_USER_PASSWORD"
  loginPat <- optionalEnv "NUTHATCH_OIDC_LOGIN_PAT"
  clientId <- requireEnv "NUTHATCH_OIDC_CLIENT_ID"
  clientSecret <- requireEnv "NUTHATCH_OIDC_CLIENT_SECRET"
  redirectUri <- requireEnv "NUTHATCH_OIDC_REDIRECT_URI"
  pure ProbeContext
    { probeClient = client
    , probeBase = "https://nuthatch-" <> lane <> ":8080"
    , probeOidcIssuer = issuer
    , probeOidcProvider = provider
    , probeOidcPassword = password
    , probeOidcLoginPat = loginPat
    , probeOidcClientId = clientId
    , probeOidcClientSecret = clientSecret
    , probeOidcRedirectUri = redirectUri
    }

login :: ProbeContext -> Bool -> String -> String -> IO Session
login context verifyBoundary user suffix = do
  cookies <- newCookieSession
  idpCookies <- newCookieSession
  started <-
    retryTransport 20 $
      getWith
        (probeClient context)
        defaultRequestOptions { optionCookies = Just cookies }
        (probeBase context <> "/login")
  assertEqual "login redirect status" 303 (probeStatus started)
  authorizationUrl <-
    maybe (fail "login response omitted Location") (pure . B8.unpack) $
      headerValue "location" started
  requireCookieAttributes
    "OIDC transaction"
    ["secure", "httponly", "samesite=lax", "path=/oidc/callback"]
    (findCookie "nuthatch_oidc" started)
  state <- queryValue "state" authorizationUrl
  if verifyBoundary
    then verifyTransactionBoundary context cookies authorizationUrl state
    else pure ()
  callback <- completeAuthorization context idpCookies user suffix authorizationUrl
  completed <- getWith
    (probeClient context)
    defaultRequestOptions { optionCookies = Just cookies }
    callback
  assertEqual "OIDC callback status" 303 (probeStatus completed)
  assertEqual "OIDC callback location"
    (Just "/app")
    (B8.unpack <$> headerValue "location" completed)
  requireCookieAttributes
    "session"
    ["secure", "httponly", "samesite=lax", "path=/;"]
    (findCookie "nuthatch_session" completed)
  requireCookieAttributes
    "CSRF"
    ["secure", "samesite=lax", "path=/;"]
    (findCookie "nuthatch_csrf" completed)
  assert "CSRF cookie must be readable by the browser" $
    not ("httponly" `containsLower` findCookie "nuthatch_csrf" completed)
  requireCookieAttributes
    "expired OIDC transaction"
    ["max-age=0", "path=/oidc/callback;"]
    (findCookie "nuthatch_oidc" completed)
  let responseCookies = headerValues "set-cookie" completed
  assert "application cookies must be host-only" $
    all (not . containsLower "domain=") responseCookies
  csrf <- cookieValue cookies "nuthatch_csrf"
  pure Session
    { sessionCookies = cookies
    , sessionCsrf = csrf
    }

verifyPkceBoundary :: ProbeContext -> String -> IO ()
verifyPkceBoundary context user = do
  cookies <- newCookieSession
  idpCookies <- newCookieSession
  started <- getWith
    (probeClient context)
    defaultRequestOptions { optionCookies = Just cookies }
    (probeBase context <> "/login")
  authorizationUrl <-
    maybe (fail "PKCE login omitted Location") (pure . B8.unpack) $
      headerValue "location" started
  callback <- completeAuthorization
    context
    idpCookies
    user
    "pkce"
    authorizationUrl
  code <- queryValue "code" callback
  discoveryResponse <- get
    (probeClient context)
    (probeOidcIssuer context <> "/.well-known/openid-configuration")
  assertEqual "OIDC discovery status" 200 (probeStatus discoveryResponse)
  OidcDiscovery tokenEndpoint <-
    decodeBody "OIDC discovery" discoveryResponse
  rejected <- postWith
    (probeClient context)
    defaultRequestOptions
      { optionBasicAuth = Just
          ( B8.pack (probeOidcClientId context)
          , B8.pack (probeOidcClientSecret context)
          )
      , optionBody = FormBody
          [ ("grant_type", "authorization_code")
          , ("code", B8.pack code)
          , ("redirect_uri", B8.pack (probeOidcRedirectUri context))
          , ( "code_verifier"
            , "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"
            )
          ]
      }
    tokenEndpoint
  assert "provider accepted an invalid PKCE verifier" $
    probeStatus rejected >= 400

verifyTransactionBoundary
  :: ProbeContext
  -> CookieSession
  -> String
  -> String
  -> IO ()
verifyTransactionBoundary context cookies authorizationUrl state = do
  wrongRedirect <- get (probeClient context) (replaceRedirectUri authorizationUrl)
  assert "provider accepted an unregistered redirect URI" $
    probeStatus wrongRedirect >= 400
  invalidState <- getWith
    (probeClient context)
    defaultRequestOptions { optionCookies = Just cookies }
    (probeBase context
      <> "/oidc/callback?code=unused&state=invalid-"
      <> state)
  assert "application accepted an invalid OIDC state" $
    probeStatus invalidState >= 400
  transaction <- cookieValue cookies "nuthatch_oidc"
  malformed <- getWith
    (probeClient context)
    defaultRequestOptions
      { optionHeaders = [("Cookie", "nuthatch_oidc=malformed")] }
    (probeBase context
      <> "/oidc/callback?code=unused&state="
      <> state)
  assertEqual "malformed transaction cookie" 400 (probeStatus malformed)
  let tampered = BS.cons 120 (BS.drop 1 transaction)
  tamperedResponse <- getWith
    (probeClient context)
    defaultRequestOptions
      { optionHeaders = [("Cookie", "nuthatch_oidc=" <> tampered)] }
    (probeBase context
      <> "/oidc/callback?code=unused&state="
      <> state)
  assertEqual "tampered transaction cookie" 400 (probeStatus tamperedResponse)

completeAuthorization
  :: ProbeContext
  -> CookieSession
  -> String
  -> String
  -> String
  -> IO String
completeAuthorization context idpCookies user suffix authorizationUrl =
  case probeOidcProvider context of
    "keycloak" ->
      completeKeycloak context idpCookies user authorizationUrl
    "zitadel" ->
      completeZitadel context user suffix authorizationUrl
    provider ->
      fail ("unsupported OIDC provider: " <> provider)

completeKeycloak
  :: ProbeContext
  -> CookieSession
  -> String
  -> String
  -> IO String
completeKeycloak context idpCookies user authorizationUrl = do
  form <- getWith
    (probeClient context)
    defaultRequestOptions { optionCookies = Just idpCookies }
    authorizationUrl
  assertEqual "Keycloak authorization form" 200 (probeStatus form)
  action <- htmlAttribute "action" (B8.unpack (probeBody form))
  submitted <- postWith
    (probeClient context)
    defaultRequestOptions
      { optionCookies = Just idpCookies
      , optionBody = FormBody
          [ ("username", B8.pack user)
          , ("password", B8.pack (probeOidcPassword context))
          , ("credentialId", "")
          ]
      }
    action
  maybe
    (fail "Keycloak login omitted callback Location")
    (pure . B8.unpack)
    (headerValue "location" submitted)

completeZitadel
  :: ProbeContext
  -> String
  -> String
  -> String
  -> IO String
completeZitadel context user _ authorizationUrl = do
  pat <- maybe (fail "ZITADEL login PAT is missing") pure $
    probeOidcLoginPat context
  started <- get (probeClient context) authorizationUrl
  location <- maybe
    (fail "ZITADEL login omitted Location")
    (pure . B8.unpack)
    (headerValue "location" started)
  authRequest <- queryValue "authRequest" location
  created <- postWith
    (probeClient context)
    defaultRequestOptions
      { optionHeaders =
          [ ("Authorization", "Bearer " <> B8.pack pat)
          , ("Content-Type", "application/json")
          ]
      , optionBody = ByteBody . B8.pack $
          "{\"checks\":{\"user\":{\"userId\":\""
      <> user
      <> "\"},\"password\":{\"password\":\""
      <> probeOidcPassword context
      <> "\"}},\"lifetime\":\"300s\"}"
      }
    (probeOidcIssuer context <> "/v2/sessions")
  assert "ZITADEL session creation failed" $
    probeStatus created >= 200 && probeStatus created < 300
  ZitadelSession {..} <- decodeBody "ZITADEL session" created
  finalized <- postWith
    (probeClient context)
    defaultRequestOptions
      { optionHeaders =
          [ ("Authorization", "Bearer " <> B8.pack pat)
          , ("Content-Type", "application/json")
          ]
      , optionBody = ByteBody . B8.pack $
          "{\"session\":{\"sessionId\":\""
      <> zitadelSessionId
      <> "\",\"sessionToken\":\""
      <> zitadelSessionToken
      <> "\"}}"
      }
    (probeOidcIssuer context
      <> "/v2/oidc/auth_requests/"
      <> authRequest)
  assertEqual "ZITADEL authorization finalization" 200 (probeStatus finalized)
  ZitadelCallback callback <- decodeBody "ZITADEL callback" finalized
  pure callback

decodeBody :: FromJSON value => String -> ProbeResponse -> IO value
decodeBody label response =
  either (fail . ((label <> ": ") <>)) pure $
    eitherDecodeStrict' (probeBody response)

findCookie :: BS.ByteString -> ProbeResponse -> BS.ByteString
findCookie name response =
  fromMaybe BS.empty $
    find (containsLower (lower name <> "=")) (headerValues "set-cookie" response)

requireCookieAttributes
  :: String
  -> [BS.ByteString]
  -> BS.ByteString
  -> IO ()
requireCookieAttributes label attributes cookie = do
  assert (label <> " cookie is missing") (not (BS.null cookie))
  mapM_
    (\attribute ->
      assert (label <> " cookie omitted " <> B8.unpack attribute) $
        attribute `containsLower` cookie)
    attributes

containsLower :: BS.ByteString -> BS.ByteString -> Bool
containsLower needle haystack =
  lower needle `BS.isInfixOf` lower haystack

lower :: BS.ByteString -> BS.ByteString
lower = B8.map toLower

queryValue :: String -> String -> IO String
queryValue name url =
  case drop 1 (dropWhile (/= '?') url) of
    "" -> fail ("URL has no query: " <> url)
    query ->
      case find ((== name) . fst) (map splitPair (splitOn '&' query)) of
        Just (_, value) | not (null value) -> pure value
        _ -> fail ("URL query omitted " <> name)

splitPair :: String -> (String, String)
splitPair value =
  case break (== '=') value of
    (name, '=' : rest) -> (name, rest)
    (name, _) -> (name, "")

splitOn :: Char -> String -> [String]
splitOn delimiter value =
  case break (== delimiter) value of
    (prefix, []) -> [prefix]
    (prefix, _ : suffix) -> prefix : splitOn delimiter suffix

replaceRedirectUri :: String -> String
replaceRedirectUri url =
  let marker = "redirect_uri="
      (prefix, rest) = breakOn marker url
  in case rest of
      "" -> url
      _ ->
        let afterName = drop (length marker) rest
            suffix = dropWhile (/= '&') afterName
        in prefix
          <> marker
          <> "https%3A%2F%2Finvalid.example%2Fcallback"
          <> suffix

htmlAttribute :: String -> String -> IO String
htmlAttribute name html =
  let marker = name <> "=\""
      (_, rest) = breakOn marker html
  in case rest of
      "" -> fail ("HTML omitted " <> name)
      _ ->
        pure
          ( replaceAll "&amp;" "&"
          $ takeWhile (/= '"')
          $ drop (length marker) rest
          )

breakOn :: String -> String -> (String, String)
breakOn needle = go []
  where
    go prefix remaining
      | needle `isPrefixOf` remaining = (reverse prefix, remaining)
    go prefix [] = (reverse prefix, "")
    go prefix (next : remaining) = go (next : prefix) remaining

isPrefixOf :: String -> String -> Bool
isPrefixOf prefix value =
  take (length prefix) value == prefix

replaceAll :: String -> String -> String -> String
replaceAll needle replacement value =
  case breakOn needle value of
    (prefix, "") -> prefix
    (prefix, rest) ->
      prefix
      <> replacement
      <> replaceAll needle replacement (drop (length needle) rest)

optionalEnv :: String -> IO (Maybe String)
optionalEnv = lookupEnv

retryTransport :: Int -> IO value -> IO value
retryTransport attempts action = do
  result <- try action
  case result of
    Right value -> pure value
    Left exception
      | attempts > 1 -> do
          threadDelay 100000
          retryTransport (attempts - 1) action
      | otherwise -> fail (show (exception :: SomeException))