-- 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 Test.Auth where

import API.Brig
import API.BrigInternal
import API.Common
import API.GalleyInternal
import qualified API.Nginz as Nginz
import API.Spar
import qualified Data.Aeson as A
import qualified Data.ByteString.Char8 as BSChar8
import qualified Data.Text as T
import SetupHelpers
import Testlib.JSON
import Testlib.Prelude
import Text.Read
import UnliftIO.Async
import UnliftIO.Concurrent

-- Happy flow: login yields a valid zauth token.
--
-- See also: 'testBearerToken'
testBearerToken2 :: (HasCallStack) => App ()
testBearerToken2 = do
  alice <- randomUser OwnDomain def
  email <- asString $ alice %. "email"
  loginResp <- login alice email defPassword >>= getJSON 200
  token <- asString $ loginResp %. "access_token"

  req <-
    rawBaseRequest alice Nginz Versioned "/self"
      <&> addHeader "Authorization" ("Bearer " <> token)
  submit "GET" req `bindResponse` \resp -> do
    resp.status `shouldMatchInt` 200
    resp.json %. "email" `shouldMatch` email

-- Happy flow (zauth token encoded in AWS4_HMAC_SHA256)
--
-- See also: 'testAWS4_HMAC_SHA256_token'
testAWS4_HMAC_SHA256_token2 :: (HasCallStack) => App ()
testAWS4_HMAC_SHA256_token2 = do
  alice <- randomUser OwnDomain def
  email <- asString $ alice %. "email"
  loginResp <- login alice email defPassword >>= getJSON 200
  token <- asString $ loginResp %. "access_token"

  let testCases =
        [ (True, "AWS4-HMAC-SHA256 Credential=" <> token <> ", foo=bar"),
          (True, "AWS4-HMAC-SHA256 Credential=" <> token),
          (True, "AWS4-HMAC-SHA256 foo=bar, Credential=" <> token),
          (True, "AWS4-HMAC-SHA256 foo=bar, Credential=" <> token <> ", baz=qux"),
          (True, "AWS4-HMAC-SHA256 foo=bar,Credential=" <> token <> ",baz=qux"),
          (False, "AWS4-HMAC-SHA256 Credential=badtoken")
        ]

  for_ testCases $ \(good, header) -> do
    req <-
      rawBaseRequest alice Nginz Versioned "/self"
        <&> addHeader "Authorization" header
    submit "GET" req `bindResponse` \resp -> do
      if good
        then do
          resp.status `shouldMatchInt` 200
          resp.json %. "email" `shouldMatch` email
        else do
          resp.status `shouldMatchInt` 401

-- The testLimitRetries test conforms to the following testing standards:
-- @SF.Channel @TSFI.RESTfulAPI @TSFI.NTP @S2
--
-- The following test tests the login retries. It checks that a user can make
-- only a prespecified number of attempts to log in with an invalid password,
-- after which the user is unable to try again for a configured amount of time.
-- After the configured amount of time has passed, the test asserts the user can
-- successfully log in again. Furthermore, the test asserts that another
-- unrelated user can successfully log-in in parallel to the failed attempts of
-- the aforementioned user.
testLimitRetries :: (HasCallStack) => App ()
testLimitRetries = do
  let retryLimit = 5
      timeout = 5
  withModifiedBackend
    def
      { brigCfg =
          -- Set a small timeout to make this test fast
          setField @_ @Int "optSettings.setLimitFailedLogins.timeout" timeout
            >=> setField @_ @Int "optSettings.setLimitFailedLogins.retryLimit" retryLimit
            -- Disable password hashing rate limiting, so we can login many times without making this test slow
            >=> setField @_ @Int "optSettings.setPasswordHashingRateLimit.userLimit.inverseRate" 0
      }
    $ \domain -> do
      alice <- randomUser domain def
      aliceEmail <- asString $ alice %. "email"

      bob <- randomUser domain def
      bobEmail <- asString $ bob %. "email"

      -- Alice tries to login 5 times with wrong password
      forM_ [1 .. retryLimit] $ \_ ->
        login domain aliceEmail "wrong-password" `bindResponse` \resp -> do
          resp.status `shouldMatchInt` 403
          resp.json %. "label" `shouldMatch` "invalid-credentials"

      -- Now alice cannot login even with correct password
      retryAfter <-
        login domain aliceEmail defPassword `bindResponse` \resp -> do
          resp.status `shouldMatchInt` 403
          resp.json %. "label" `shouldMatch` "client-error"
          let Just retryAfter = readMaybe . BSChar8.unpack =<< lookup (fromString "Retry-After") resp.headers
          (retryAfter <= timeout) `shouldMatch` True
          pure retryAfter

      -- Bob can still login
      login domain bobEmail defPassword `bindResponse` \resp -> do
        resp.status `shouldMatchInt` 200

      -- Waiting 2s less than retryAfter should still cause a failure
      threadDelay ((retryAfter - 2) * 1_000_000)
      login domain aliceEmail defPassword `bindResponse` \resp -> do
        resp.status `shouldMatchInt` 403
        resp.json %. "label" `shouldMatch` "client-error"
        let Just retryAfter2 = readMaybe . BSChar8.unpack =<< lookup (fromString "Retry-After") resp.headers
        -- This should be about 2 seconds or slightly less because we didn't
        -- wait long enough. This also asserts that the throttling doesn't get
        -- reset by making another call
        (retryAfter2 <= (2 :: Int)) `shouldMatch` True

      -- Waiting 2 more seconds should make the login succeed
      threadDelay 2_000_000
      login domain aliceEmail defPassword `bindResponse` \resp -> do
        resp.status `shouldMatchInt` 200

-- @END

-- The testTooManyCookies test conforms to the following testing standards:
-- @SF.Provisioning @TSFI.RESTfulAPI @S2
--
-- The test asserts that there is an upper limit for the number of user cookies
-- per cookie type. It does that by concurrently attempting to create more
-- persistent and session cookies than the configured maximum.
-- Creation of new cookies beyond the limit causes deletion of the
-- oldest cookies.
testTooManyCookies :: (HasCallStack) => App ()
testTooManyCookies = do
  let cookieLimit = 5
  withModifiedBackend
    def
      { brigCfg =
          -- Disable password hashing rate limiting, so we can login many times without making this test slow
          setField @_ @Int "optSettings.setPasswordHashingRateLimit.userLimit.inverseRate" 0
            -- Disable cookie throttling so this test is not slow
            >=> setField @_ @Int "optSettings.setUserCookieThrottle.retryAfter" 0
            >=> setField @_ @Int "optSettings.setUserCookieLimit" cookieLimit
      }
    $ \domain -> do
      alice <- randomUser domain def
      aliceEmail <- asString $ alice %. "email"

      let testCookieLimit label = do
            let loginFn = if label == "persistent" then login else loginWithSessionCookie
            (deletedCookie1 : deletedCookie2 : validCookies) <-
              replicateM (cookieLimit + 2)
                $ do
                  -- This threadDelay is required to get around problems caused
                  -- by: https://wearezeta.atlassian.net/browse/WPB-15446
                  threadDelay 1_000_000
                  loginFn domain aliceEmail defPassword
                    `bindResponse` \resp -> do
                      resp.status `shouldMatchInt` 200
                      pure . fromJust $ getCookie "zuid" resp
            addFailureContext ("deletedCookie1: " <> deletedCookie1 <> "\ndeletedCookie2: " <> deletedCookie2 <> "\nvalidCookies:\n" <> unlines validCookies) $ do
              forM_ [deletedCookie1, deletedCookie2] $ \deletedCookie -> do
                renewToken alice deletedCookie `bindResponse` \resp ->
                  resp.status `shouldMatchInt` 403
              forM_ validCookies $ \validCookie ->
                renewToken alice validCookie `bindResponse` \resp ->
                  resp.status `shouldMatchInt` 200
      concurrently_ (testCookieLimit "persistent") (testCookieLimit "session")

-- @END

-- The testInvalidCookie test conforms to the following testing standards:
-- @SF.Provisioning @TSFI.RESTfulAPI @TSFI.NTP @S2
--
-- Test that invalid and expired tokens do not work.
testInvalidCookie :: (HasCallStack) => App ()
testInvalidCookie = do
  let cookieTimeout = 2
  withModifiedBackend
    def
      { brigCfg =
          setField @_ @Int "zauth.authSettings.userTokenTimeout" cookieTimeout
            >=> setField @_ @Int "zauth.authSettings.legalHoldUserTokenTimeout" cookieTimeout
      }
    $ \domain -> do
      Nginz.access domain "zuid=xxx" `bindResponse` \resp -> do
        resp.status `shouldMatchInt` 403
        resp.json %. "label" `shouldMatch` "client-error"
        msg <- asString $ resp.json %. "message"
        msg `shouldContain` "Invalid zauth token"

      (owner, tid, [alice]) <- createTeam domain 2
      aliceEmail <- asString $ alice %. "email"
      aliceId <- asString $ alice %. "qualified_id.id"
      userCookie <-
        login domain aliceEmail defPassword `bindResponse` \resp -> do
          resp.status `shouldMatchInt` 200
          pure . fromJust $ getCookie "zuid" resp

      legalholdWhitelistTeam tid owner >>= assertSuccess

      lhCookie <-
        legalholdLogin domain aliceId defPassword `bindResponse` \resp -> do
          resp.status `shouldMatchInt` 200
          pure . fromJust $ getCookie "zuid" resp

      -- Wait for both cookies to expire
      -- In order to reduce flakiness, the delay has been increased by 2 seconds
      -- after an investigation revealed that check for cookie expiration is somewhat lenient due to
      -- the way the time is calculated in the backend.
      -- See the interpreter of Now which is implemented using `Control.AutoUpdate` which defaults to an update frequency of 1 sec.
      -- Furthermore the timestamp is then again rounded down to the nearest second before it is compared to the cookie expiration time.
      threadDelay $ (cookieTimeout + 2) * 1_000_000
      -- Assert that the cookies are considered expired
      for_ [userCookie, lhCookie] $ \cookie ->
        Nginz.access domain ("zuid=" <> cookie) `bindResponse` \resp -> do
          resp.status `shouldMatchInt` 403
          resp.json %. "label" `shouldMatch` "invalid-credentials"
          resp.json %. "message" `shouldMatch` "Zauth token expired"

-- @END

testSetCookieLabelOnSsoFlow :: (HasCallStack) => App ()
testSetCookieLabelOnSsoFlow = do
  let sharedDeviceMarker = "shared-device"
  domain <- make OwnDomain
  (owner, tid, _) <- createTeam OwnDomain 1
  void $ setTeamFeatureStatus owner tid "sso" "enabled"
  void $ setTeamFeatureStatus owner tid "validateSAMLemails" "enabled"
  idp@(samlId, _) <- do
    (resp, (meta, creds)) <- registerTestIdPWithMetaWithPrivateCreds owner
    resp.status `shouldMatchInt` 201
    (,(meta, creds)) <$> (resp.json %. "id" >>= asString)
  scimToken <- do
    let p = def {name = Just "my-idp", idp = Just samlId}
    createScimToken owner p >>= getJSON 200 >>= (%. "token") >>= asString

  [(email, user), (email2, user2)] <- replicateM 2 do
    scimUser <- randomScimUser
    email <- scimUser %. "externalId" >>= asString
    uid <- bindResponse (createScimUser owner scimToken scimUser) $ \resp -> do
      resp.status `shouldMatchInt` 201
      resp.json %. "id" >>= asString
    let user = object ["id" .= uid, "qualified_id" .= object ["id" .= uid, "domain" .= domain]]
    pure (email, user)

  let getJust = maybe (assertFailure "Expected a cookie, but got no cookie") pure

  -- labeled cookie
  cookie1 <- getCookieWithSaml tid email idp (Just sharedDeviceMarker) >>= getJust
  checkCookies user [] [cookie1] [sharedDeviceMarker]

  withWebSocket user \userWs -> do
    -- same label
    cookie2 <- getCookieWithSaml tid email idp (Just sharedDeviceMarker) >>= getJust
    checkCookies user [cookie1] [cookie2] [sharedDeviceMarker]

    void $ awaitMatch isUserSessionRefreshSuggested userWs

    -- other label
    let otherLabel = "funky device"
    cookie3 <- getCookieWithSaml tid email idp (Just otherLabel) >>= getJust

    -- unlabeled cookie
    cookie4 <- getCookieWithSaml tid email idp Nothing >>= getJust
    checkCookies user [cookie1] [cookie2, cookie3, cookie4] [sharedDeviceMarker, otherLabel]

    -- other user with same label
    cookie5 <- getCookieWithSaml tid email2 idp (Just sharedDeviceMarker) >>= getJust
    checkCookies user [cookie1] [cookie2, cookie3, cookie4] [sharedDeviceMarker, otherLabel]
    checkCookies user2 [] [cookie5] [sharedDeviceMarker]

    -- another login with the same label
    cookie6 <- getCookieWithSaml tid email idp (Just sharedDeviceMarker) >>= getJust
    checkCookies user [cookie1, cookie2] [cookie3, cookie4, cookie6] [sharedDeviceMarker, otherLabel]

    void $ awaitMatch isUserSessionRefreshSuggested userWs
  where
    isUserSessionRefreshSuggested :: (HasCallStack, MakesValue a) => a -> App Bool
    isUserSessionRefreshSuggested n = fieldEquals n "payload.0.type" "user.session-refresh-suggested"

    checkCookies user expectedInvalid expectedValid expectedLabels = do
      checkCookiesValid user expectedValid expectedLabels
      checkCookiesInvalid expectedInvalid

    checkCookiesInvalid cookies =
      for_ cookies \cookie ->
        Nginz.access OwnDomain ("zuid=" <> cookie) `bindResponse` \resp -> do
          resp.status `shouldMatchInt` 403
          resp.json %. "label" `shouldMatch` "invalid-credentials"
          resp.json %. "message" `shouldMatch` "Invalid zauth token"

    checkCookiesValid user expectedCookies expectedLabels = do
      for_ expectedCookies $ \cookie -> Nginz.access OwnDomain ("zuid=" <> cookie) >>= assertSuccess
      getCookies user [] `bindResponse` \resp -> do
        resp.status `shouldMatchInt` 200
        cookies <- resp.json %. "cookies" >>= asList
        length cookies `shouldMatchInt` length expectedCookies
        cookieLabels <- traverse (\c -> c %. "label") cookies
        let expected = sort $ (A.String . T.pack <$> expectedLabels)
        let actual = sort $ filter ((/=) A.Null) cookieLabels
        actual `shouldMatch` expected
