-- This file is part of the Wire Server implementation.
--
-- Copyright (C) 2026 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.TOTP where

import qualified API.Brig as Brig
import API.Common (defPassword)
import qualified Data.ByteString.Char8 as BSChar8
import qualified Data.CaseInsensitive as CI
import qualified Network.HTTP.Types.URI as HttpUri
import qualified Network.URI as NetUri
import SetupHelpers
import Testlib.Prelude

--------------------------------------------------------------------------------
-- mock bodies

mockEnrolBody :: Value
mockEnrolBody :: Value
mockEnrolBody =
  [Pair] -> Value
object
    [ String
"name" String -> String -> Pair
forall a. ToJSON a => String -> a -> Pair
.= (String
"work phone" :: String),
      String
"client_time" String -> String -> Pair
forall a. ToJSON a => String -> a -> Pair
.= (String
"2026-02-03T11:17:05.021Z" :: String),
      String
"password" String -> String -> Pair
forall a. ToJSON a => String -> a -> Pair
.= String
defPassword
    ]

mockActivateBody :: Value
mockActivateBody :: Value
mockActivateBody =
  [Pair] -> Value
object
    [ String
"code" String -> String -> Pair
forall a. ToJSON a => String -> a -> Pair
.= (String
"014921" :: String),
      String
"client_time" String -> String -> Pair
forall a. ToJSON a => String -> a -> Pair
.= (String
"2026-02-03T11:18:44.903Z" :: String)
    ]

mockRenameBody :: Value
mockRenameBody :: Value
mockRenameBody = [Pair] -> Value
object [String
"name" String -> String -> Pair
forall a. ToJSON a => String -> a -> Pair
.= (String
"private phone" :: String)]

mockDeleteBody :: Value
mockDeleteBody :: Value
mockDeleteBody = [Pair] -> Value
object [String
"password" String -> String -> Pair
forall a. ToJSON a => String -> a -> Pair
.= String
defPassword]

--------------------------------------------------------------------------------
-- POST /self/totp

testTotpEnrol :: (HasCallStack) => App ()
testTotpEnrol :: HasCallStack => App ()
testTotpEnrol = do
  alice <- Domain -> CreateUser -> App Value
forall domain.
(HasCallStack, MakesValue domain) =>
domain -> CreateUser -> App Value
randomUser Domain
OwnDomain CreateUser
forall a. Default a => a
def
  bindResponse (Brig.enrolTotp alice mockEnrolBody) $ \Response
resp -> do
    Response
resp.status Int -> Int -> App ()
forall a. (MakesValue a, HasCallStack) => a -> Int -> App ()
`shouldMatchInt` Int
201

    -- the one response in the set that carries a secret must not be cached
    let cacheControl :: Maybe String
cacheControl = (Char -> Char) -> String -> String
forall a b. (a -> b) -> [a] -> [b]
map Char -> Char
toLower (String -> String)
-> (ByteString -> String) -> ByteString -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> String
BSChar8.unpack (ByteString -> String) -> Maybe ByteString -> Maybe String
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> Response -> Maybe ByteString
responseHeader String
"Cache-Control" Response
resp
    HasCallStack => String -> Bool -> App ()
String -> Bool -> App ()
assertBool
      (String
"expected Cache-Control: no-store, got " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Maybe String -> String
forall a. Show a => a -> String
show Maybe String
cacheControl)
      (Maybe String
cacheControl Maybe String -> Maybe String -> Bool
forall a. Eq a => a -> a -> Bool
== String -> Maybe String
forall a. a -> Maybe a
Just String
"no-store")

    App String -> App ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (App String -> App ()) -> App String -> App ()
forall a b. (a -> b) -> a -> b
$ App Value -> App String
forall a. (HasCallStack, MakesValue a) => a -> App String
asString (Response
resp.json Maybe Value -> String -> App Value
forall a. (HasCallStack, MakesValue a) => a -> String -> App Value
%. String
"credential_id")
    App String -> App ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (App String -> App ()) -> App String -> App ()
forall a b. (a -> b) -> a -> b
$ App Value -> App String
forall a. (HasCallStack, MakesValue a) => a -> App String
asString (Response
resp.json Maybe Value -> String -> App Value
forall a. (HasCallStack, MakesValue a) => a -> String -> App Value
%. String
"secret")
    App String -> App ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (App String -> App ()) -> App String -> App ()
forall a b. (a -> b) -> a -> b
$ App Value -> App String
forall a. (HasCallStack, MakesValue a) => a -> App String
asString (Response
resp.json Maybe Value -> String -> App Value
forall a. (HasCallStack, MakesValue a) => a -> String -> App Value
%. String
"uri")
    alg <- App Value -> App String
forall a. (HasCallStack, MakesValue a) => a -> App String
asString (Response
resp.json Maybe Value -> String -> App Value
forall a. (HasCallStack, MakesValue a) => a -> String -> App Value
%. String
"algorithm")
    digits <- asInt (resp.json %. "digits")
    period <- asInt (resp.json %. "period")
    void $ asString (resp.json %. "server_time")
    void $ asString (resp.json %. "expires_at")

    -- the mock pins these rather than drawing them at random, so that a client
    -- that hardcodes the defaults does not look broken against it
    alg `shouldMatch` "sha1"
    digits `shouldMatchInt` 6
    period `shouldMatchInt` 30

-- | secret and uri field in response match.
testTotpEnrolUriMatchesSecret :: (HasCallStack) => App ()
testTotpEnrolUriMatchesSecret :: HasCallStack => App ()
testTotpEnrolUriMatchesSecret = do
  alice <- Domain -> CreateUser -> App Value
forall domain.
(HasCallStack, MakesValue domain) =>
domain -> CreateUser -> App Value
randomUser Domain
OwnDomain CreateUser
forall a. Default a => a
def
  bindResponse (Brig.enrolTotp alice mockEnrolBody) $ \Response
resp -> do
    Response
resp.status Int -> Int -> App ()
forall a. (MakesValue a, HasCallStack) => a -> Int -> App ()
`shouldMatchInt` Int
201
    secret <- App Value -> App String
forall a. (HasCallStack, MakesValue a) => a -> App String
asString (Response
resp.json Maybe Value -> String -> App Value
forall a. (HasCallStack, MakesValue a) => a -> String -> App Value
%. String
"secret")
    uriText <- asString (resp.json %. "uri")

    uri <- case NetUri.parseURI uriText of
      Maybe URI
Nothing -> String -> App URI
forall a. HasCallStack => String -> App a
assertFailure (String
"uri does not parse as a URI: " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
uriText)
      Just URI
u -> URI -> App URI
forall a. a -> App a
forall (f :: * -> *) a. Applicative f => a -> f a
pure URI
u
    NetUri.uriScheme uri `shouldMatch` "otpauth:"
    case NetUri.uriAuthority uri of
      Just URIAuth
auth -> URIAuth -> String
NetUri.uriRegName URIAuth
auth String -> String -> App ()
forall a b.
(MakesValue a, MakesValue b, HasCallStack) =>
a -> b -> App ()
`shouldMatch` String
"totp"
      Maybe URIAuth
Nothing -> String -> App ()
forall a. HasCallStack => String -> App a
assertFailure (String
"uri has no authority: " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
uriText)

    let params = ByteString -> Query
HttpUri.parseQuery (String -> ByteString
BSChar8.pack (URI -> String
NetUri.uriQuery URI
uri))
        param String
k =
          App String
-> (ByteString -> App String) -> Maybe ByteString -> App String
forall b a. b -> (a -> b) -> Maybe a -> b
maybe
            (String -> App String
forall a. HasCallStack => String -> App a
assertFailure (String
"uri has no " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
k String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
" parameter: " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
uriText))
            (String -> App String
forall a. a -> App a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (String -> App String)
-> (ByteString -> String) -> ByteString -> App String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> String
BSChar8.unpack)
            (Maybe (Maybe ByteString) -> Maybe ByteString
forall (m :: * -> *) a. Monad m => m (m a) -> m a
join (ByteString -> Query -> Maybe (Maybe ByteString)
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup (String -> ByteString
BSChar8.pack String
k) Query
params))
    -- this is the assertion the whole test exists for
    param "secret" `shouldMatch` secret
    param "issuer" `shouldMatch` "Wire"
    param "algorithm" `shouldMatch` "SHA1"
    param "digits" `shouldMatch` "6"
    param "period" `shouldMatch` "30"

--------------------------------------------------------------------------------
-- the other four routes

testTotpActivate :: (HasCallStack) => App ()
testTotpActivate :: HasCallStack => App ()
testTotpActivate = do
  alice <- Domain -> CreateUser -> App Value
forall domain.
(HasCallStack, MakesValue domain) =>
domain -> CreateUser -> App Value
randomUser Domain
OwnDomain CreateUser
forall a. Default a => a
def
  cid <- someCredentialId alice
  bindResponse (Brig.activateTotp alice cid mockActivateBody) $ \Response
resp -> do
    Response
resp.status Int -> Int -> App ()
forall a. (MakesValue a, HasCallStack) => a -> Int -> App ()
`shouldMatchInt` Int
200
    codes <- Response
resp.json Maybe Value -> String -> App Value
forall a. (HasCallStack, MakesValue a) => a -> String -> App Value
%. String
"codes" App Value -> (App Value -> App [Value]) -> App [Value]
forall a b. a -> (a -> b) -> b
& App Value -> App [Value]
forall a. (HasCallStack, MakesValue a) => a -> App [Value]
asList
    for_ codes $ \Value
c -> do
      code <- Value -> App String
forall a. (HasCallStack, MakesValue a) => a -> App String
asString Value
c
      assertBool
        ("recovery code has the wrong shape: " <> code)
        (isRecoveryCodeShaped code)

testTotpList :: (HasCallStack) => App ()
testTotpList :: HasCallStack => App ()
testTotpList = do
  alice <- Domain -> CreateUser -> App Value
forall domain.
(HasCallStack, MakesValue domain) =>
domain -> CreateUser -> App Value
randomUser Domain
OwnDomain CreateUser
forall a. Default a => a
def
  bindResponse (Brig.listTotp alice) $ \Response
resp -> do
    Response
resp.status Int -> Int -> App ()
forall a. (MakesValue a, HasCallStack) => a -> Int -> App ()
`shouldMatchInt` Int
200
    creds <- Maybe Value -> App [Value]
forall a. (HasCallStack, MakesValue a) => a -> App [Value]
asList Response
resp.json
    for_ creds $ \Value
c -> do
      App String -> App ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (App String -> App ()) -> App String -> App ()
forall a b. (a -> b) -> a -> b
$ App Value -> App String
forall a. (HasCallStack, MakesValue a) => a -> App String
asString (Value
c Value -> String -> App Value
forall a. (HasCallStack, MakesValue a) => a -> String -> App Value
%. String
"id")
      App String -> App ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (App String -> App ()) -> App String -> App ()
forall a b. (a -> b) -> a -> b
$ App Value -> App String
forall a. (HasCallStack, MakesValue a) => a -> App String
asString (Value
c Value -> String -> App Value
forall a. (HasCallStack, MakesValue a) => a -> String -> App Value
%. String
"name")
      App String -> App ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (App String -> App ()) -> App String -> App ()
forall a b. (a -> b) -> a -> b
$ App Value -> App String
forall a. (HasCallStack, MakesValue a) => a -> App String
asString (Value
c Value -> String -> App Value
forall a. (HasCallStack, MakesValue a) => a -> String -> App Value
%. String
"created_at")
      -- the secret and the URI are returned by enrolment and never again
      Value -> String -> App (Maybe Value)
forall a.
(HasCallStack, MakesValue a) =>
a -> String -> App (Maybe Value)
lookupField Value
c String
"secret" App (Maybe Value) -> Maybe Value -> App ()
forall a b.
(MakesValue a, MakesValue b, HasCallStack) =>
a -> b -> App ()
`shouldMatch` Maybe Value
noValue
      Value -> String -> App (Maybe Value)
forall a.
(HasCallStack, MakesValue a) =>
a -> String -> App (Maybe Value)
lookupField Value
c String
"uri" App (Maybe Value) -> Maybe Value -> App ()
forall a b.
(MakesValue a, MakesValue b, HasCallStack) =>
a -> b -> App ()
`shouldMatch` Maybe Value
noValue

testTotpRename :: (HasCallStack) => App ()
testTotpRename :: HasCallStack => App ()
testTotpRename = do
  alice <- Domain -> CreateUser -> App Value
forall domain.
(HasCallStack, MakesValue domain) =>
domain -> CreateUser -> App Value
randomUser Domain
OwnDomain CreateUser
forall a. Default a => a
def
  cid <- someCredentialId alice
  bindResponse (Brig.renameTotp alice cid mockRenameBody) $ \Response
resp ->
    Response
resp.status Int -> Int -> App ()
forall a. (MakesValue a, HasCallStack) => a -> Int -> App ()
`shouldMatchInt` Int
200

testTotpDelete :: (HasCallStack) => App ()
testTotpDelete :: HasCallStack => App ()
testTotpDelete = do
  alice <- Domain -> CreateUser -> App Value
forall domain.
(HasCallStack, MakesValue domain) =>
domain -> CreateUser -> App Value
randomUser Domain
OwnDomain CreateUser
forall a. Default a => a
def
  cid <- someCredentialId alice
  bindResponse (Brig.deleteTotp alice cid mockDeleteBody) $ \Response
resp ->
    Response
resp.status Int -> Int -> App ()
forall a. (MakesValue a, HasCallStack) => a -> Int -> App ()
`shouldMatchInt` Int
200

--------------------------------------------------------------------------------
-- version gate

testTotpNotOnV18 :: (HasCallStack) => App ()
testTotpNotOnV18 :: HasCallStack => App ()
testTotpNotOnV18 = do
  alice <- Domain -> CreateUser -> App Value
forall domain.
(HasCallStack, MakesValue domain) =>
domain -> CreateUser -> App Value
randomUser Domain
OwnDomain CreateUser
forall a. Default a => a
def
  let routes =
        [ (String
"POST", String
"/self/totp"),
          (String
"GET", String
"/self/totp"),
          (String
"POST", String
"/self/totp/" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
credentialIdPlaceholder String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
"/activate"),
          (String
"PUT", String
"/self/totp/" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
credentialIdPlaceholder),
          (String
"DELETE", String
"/self/totp/" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
credentialIdPlaceholder)
        ]
  for_ routes $ \(String
method, String
path) -> do
    req <- Value -> Service -> Versioned -> String -> App Request
forall user.
(HasCallStack, MakesValue user) =>
user -> Service -> Versioned -> String -> App Request
baseRequest Value
alice Service
Brig (Int -> Versioned
ExplicitVersion Int
18) String
path
    bindResponse (submit method (req & addJSONObject ["name" .= ("work phone" :: String)])) $ \Response
resp ->
      Response
resp.status Int -> [Value] -> App ()
forall a b.
(MakesValue a, MakesValue b, HasCallStack) =>
a -> b -> App ()
`shouldMatchOneOf` [Scientific -> Value
Number Scientific
404, Scientific -> Value
Number Scientific
405]

--------------------------------------------------------------------------------
-- nginz

testTotpThroughNginz :: (HasCallStack) => App ()
testTotpThroughNginz :: HasCallStack => App ()
testTotpThroughNginz = do
  alice <- Domain -> CreateUser -> App Value
forall domain.
(HasCallStack, MakesValue domain) =>
domain -> CreateUser -> App Value
randomUser Domain
OwnDomain CreateUser
forall a. Default a => a
def
  email <- asString (alice %. "email")
  token <- do
    resp <- Brig.login alice email defPassword >>= getJSON 200
    asString (resp %. "access_token")

  -- authenticated: the request reaches brig
  do
    req <-
      rawBaseRequest alice Nginz Versioned "/self/totp"
        <&> addHeader "Authorization" ("Bearer " <> token)
    bindResponse (submit "GET" req) $ \Response
resp -> do
      Response
resp.status Int -> Int -> App ()
forall a. (MakesValue a, HasCallStack) => a -> Int -> App ()
`shouldMatchInt` Int
200
      App [Value] -> App ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (App [Value] -> App ()) -> App [Value] -> App ()
forall a b. (a -> b) -> a -> b
$ Maybe Value -> App [Value]
forall a. (HasCallStack, MakesValue a) => a -> App [Value]
asList Response
resp.json

  -- a bad token: nginz rejects it before brig sees anything
  do
    req <-
      rawBaseRequest alice Nginz Versioned "/self/totp"
        <&> addHeader "Authorization" "Bearer not-a-token"
    bindResponse (submit "GET" req) $ \Response
resp ->
      Response
resp.status Int -> Int -> App ()
forall a. (MakesValue a, HasCallStack) => a -> Int -> App ()
`shouldMatchInt` Int
401

  -- no credentials at all
  do
    req <- rawBaseRequest alice Nginz Versioned "/self/totp"
    bindResponse (submit "GET" req) $ \Response
resp ->
      HasCallStack => String -> Bool -> App ()
String -> Bool -> App ()
assertBool
        (String
"unauthenticated request was not rejected: " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. Show a => a -> String
show Response
resp.status)
        (Response
resp.status Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
400 Bool -> Bool -> Bool
&& Response
resp.status Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
500)

--------------------------------------------------------------------------------
-- helpers

someCredentialId :: (HasCallStack, MakesValue user) => user -> App String
someCredentialId :: forall a. (HasCallStack, MakesValue a) => a -> App String
someCredentialId user
user = do
  resp <- user -> Value -> App Response
forall user.
(HasCallStack, MakesValue user) =>
user -> Value -> App Response
Brig.enrolTotp user
user Value
mockEnrolBody App Response -> (Response -> App Value) -> App Value
forall a b. App a -> (a -> App b) -> App b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= HasCallStack => Int -> Response -> App Value
Int -> Response -> App Value
getJSON Int
201
  asString (resp %. "credential_id")

credentialIdPlaceholder :: String
credentialIdPlaceholder :: String
credentialIdPlaceholder = String
"0f2e1b9c-6c6c-4a0e-9a3f-17a2cf15a001"

responseHeader :: String -> Response -> Maybe ByteString
responseHeader :: String -> Response -> Maybe ByteString
responseHeader String
name Response
resp =
  (CI ByteString, ByteString) -> ByteString
forall a b. (a, b) -> b
snd ((CI ByteString, ByteString) -> ByteString)
-> Maybe (CI ByteString, ByteString) -> Maybe ByteString
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((CI ByteString, ByteString) -> Bool)
-> [(CI ByteString, ByteString)]
-> Maybe (CI ByteString, ByteString)
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find ((ByteString -> CI ByteString
forall s. FoldCase s => s -> CI s
CI.mk (String -> ByteString
BSChar8.pack String
name) CI ByteString -> CI ByteString -> Bool
forall a. Eq a => a -> a -> Bool
==) (CI ByteString -> Bool)
-> ((CI ByteString, ByteString) -> CI ByteString)
-> (CI ByteString, ByteString)
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (CI ByteString, ByteString) -> CI ByteString
forall a b. (a, b) -> a
fst) Response
resp.headers

-- | @xxxx-xxxx-xxxx-xxxx@ in Crockford base32 (no I, L, O, U).
isRecoveryCodeShaped :: String -> Bool
isRecoveryCodeShaped :: String -> Bool
isRecoveryCodeShaped String
s =
  ([String] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [String]
groups Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
4)
    Bool -> Bool -> Bool
&& (String -> Bool) -> [String] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all ((Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
4) (Int -> Bool) -> (String -> Int) -> String -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length) [String]
groups
    Bool -> Bool -> Bool
&& (String -> Bool) -> [String] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all ((Char -> Bool) -> String -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (Char -> String -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` String
crockford)) [String]
groups
  where
    crockford :: String
    crockford :: String
crockford = String
"0123456789ABCDEFGHJKMNPQRSTVWXYZ"

    groups :: [String]
groups = String -> [String]
splitOnDash String
s
    splitOnDash :: String -> [String]
splitOnDash String
t = case (Char -> Bool) -> String -> (String, String)
forall a. (a -> Bool) -> [a] -> ([a], [a])
break (Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== Char
'-') String
t of
      (String
g, []) -> [String
g]
      (String
g, Char
_ : String
rest) -> String
g String -> [String] -> [String]
forall a. a -> [a] -> [a]
: String -> [String]
splitOnDash String
rest