feat(monkeys): Add support for 2FA during login [WPB-6289] (#2445)

* feat(monkeys): Add support for 2FA during login [WPB-6289]

* fix

* detekt

* add workaround to reuse code

* detekt

monkeys/README.md

@@ -32,7 +32,7 @@ config file to load simply set the environment variable `CONFIG`:
 env CONFIG=example.json docker compose up
 ```
 
-*Note*: the file must be located inside the config folder.
+_Note_: the file must be located inside the config folder.
 
 If the host is not an ARM based platform, the `--platform` parameter can be omitted. If building on
 MacOs ARM computers, be sure to disable Rosetta emulation and containerd support in the settings
@@ -68,12 +68,13 @@ scalability. To do that you need to specify a command to start each monkey and a
 them. Both the `startCommand` and the `addressTemplate` are templated string the following variables
 will be replaced at runtime:
 
-* teamName: the name of the team that the user belongs to
-* email: the email of the user
-* userId: the id of the user (without the domain)
-* teamId: the id of the team that the user belongs to
-* monkeyIndex: a unique identifier within the app for each monkey (it is numeric starting from 0)
-* monkeyClientId: a unique identifier for each client within the app (it is numeric)
+- teamName: the name of the team that the user belongs to
+- email: the email of the user
+- userId: the id of the user (without the domain)
+- teamId: the id of the team that the user belongs to
+- monkeyIndex: a unique identifier within the app for each monkey (it is numeric starting from 0)
+- monkeyClientId: a unique identifier for each client within the app (it is numeric)
+- code: the 2FA code that can be reused until expired.
 
 Example:
 
@@ -100,17 +101,29 @@ needed on non x86_64 CPUs):
 }
 ```
 
+Or on swarm mode:
+
+```json
+{
+    "externalMonkey": {
+        "startCommand": "docker service create --hostname monkey-{{monkeyIndex}} --name monkey-{{monkeyIndex}}-{{teamName}} --network infinite-monkeys_monkeys --restart-condition none monkeys /opt/app/bin/monkey-server",
+        "addressTemplate": "http://monkey-{{monkeyIndex}}:8080",
+        "waitTime": 3
+    }
+}
+```
+
 The `waitTime` field is optional and determines how long (in seconds) it will wait until it proceeds to start the
 next monkey. This can be important because the next step in the app is assigning each to a user and
 the app must be ready to respond.
 
-**Note**: setting environment variables prior to the command is is not supported. Ex: 
+**Note**: setting environment variables prior to the command is is not supported. Ex:
 `INDEX={{monkeyIndex}} ./monkeys/bin/monkey-server -p 50$INDEX`. This is a limitation of the JVM's system command
 runner.
 
 ## Current Limitations (to be fixed in the future)
 
-* The application should run until it receives a `SIGINT` (Ctrl+C) signal. There should be a
-  configuration to finish the test run
-* Tests need to be implemented
-* Randomising times for action execution
+- The application should run until it receives a `SIGINT` (Ctrl+C) signal. There should be a
+    configuration to finish the test run
+- Tests need to be implemented
+- Randomising times for action execution

monkeys/schema.json

@@ -198,6 +198,11 @@
                         "description": "How many users should be created in this team",
                         "type": "integer"
                     },
+                    "2FAEnabled": {
+                        "description": "Does this server require 2FA authentication?",
+                        "type": "boolean",
+                        "default": false
+                    },
                     "dumpUsers": {
                         "description": "Should the application dump the users created into a file",
                         "type": "boolean"

monkeys/src/main/kotlin/com/wire/kalium/monkeys/Utils.kt

@@ -35,4 +35,5 @@ fun String.renderMonkeyTemplate(userData: UserData, monkeyId: MonkeyId): String
     return this.replace("{{teamName}}", userData.team.name).replace("{{email}}", userData.email)
         .replace("{{userId}}", userData.userId.value).replace("{{teamId}}", userData.team.id)
         .replace("{{monkeyIndex}}", monkeyId.index.toString()).replace("{{monkeyClientId}}", monkeyId.clientId.toString())
+        .replace("{{code}}", userData.oldCode.orEmpty())
 }

monkeys/src/main/kotlin/com/wire/kalium/monkeys/conversation/LocalMonkey.kt

@@ -64,10 +64,14 @@ class LocalMonkey(monkeyType: MonkeyType, internalId: MonkeyId) : Monkey(monkeyT
      * Logs user in and register client (if not registered)
      */
     override suspend fun login(coreLogic: CoreLogic, callback: (Monkey) -> Unit) {
-        val authScope = getAuthScope(coreLogic, this.monkeyType.userData().team.backend)
-        val email = this.monkeyType.userData().email
-        val password = this.monkeyType.userData().password
-        val loginResult = authScope.login(email, password, false)
+        val userData = this.monkeyType.userData()
+        val secondFactor = userData.request2FA()
+        val authScope = getAuthScope(coreLogic, userData.team.backend)
+        val email = userData.email
+        val password = userData.password
+        val loginResult = authScope.login(
+            userIdentifier = email, password = password, shouldPersistClient = false, secondFactorVerificationCode = secondFactor
+        )
         if (loginResult !is AuthenticationResult.Success) {
             error("User creds didn't work ($email, $password)")
         }
@@ -86,7 +90,10 @@ class LocalMonkey(monkeyType: MonkeyType, internalId: MonkeyId) : Monkey(monkeyT
         }
         val sessionScope = coreLogic.getSessionScope(loginResult.authData.userId)
         val registerClientParam = RegisterClientUseCase.RegisterClientParam(
-            password = this.monkeyType.userData().password, capabilities = emptyList(), clientType = ClientType.Temporary
+            password = userData.password,
+            capabilities = emptyList(),
+            clientType = ClientType.Temporary,
+            secondFactorVerificationCode = secondFactor
         )
         val registerResult = sessionScope.client.getOrRegister(registerClientParam)
         if (registerResult is RegisterClientResult.Failure) {

monkeys/src/main/kotlin/com/wire/kalium/monkeys/model/TestData.kt

@@ -188,6 +188,7 @@ data class BackendConfig(
     @SerialName("authUser") val authUser: String,
     @SerialName("authPassword") val authPassword: String,
     @SerialName("userCount") val userCount: ULong,
+    @SerialName("2FAEnabled") val secondFactorAuth: Boolean = false,
     @SerialName("dumpUsers") val dumpUsers: Boolean = false,
     @SerialName("presetTeam") val presetTeam: TeamConfig? = null
 )

monkeys/src/main/kotlin/com/wire/kalium/monkeys/model/TestDataImporter.kt

@@ -46,6 +46,7 @@ import kotlinx.serialization.ExperimentalSerializationApi
 import kotlinx.serialization.encodeToString
 import kotlinx.serialization.json.Json
 import kotlinx.serialization.json.JsonObject
+import kotlinx.serialization.json.JsonPrimitive
 import kotlinx.serialization.json.decodeFromStream
 import kotlinx.serialization.json.jsonArray
 import kotlinx.serialization.json.jsonObject
@@ -97,22 +98,26 @@ object TestDataImporter {
                     backendConfig.presetTeam.owner.email,
                     backendConfig.passwordForUsers,
                     UserId(backendConfig.presetTeam.owner.unqualifiedId, backendConfig.domain),
+                    null,
                     httpClient
                 )
                 backendConfig.presetTeam.users.map { user ->
                     UserData(
-                        user.email, backendConfig.passwordForUsers, UserId(user.unqualifiedId, backendConfig.domain), team
+                        user.email,
+                        backendConfig.passwordForUsers,
+                        UserId(user.unqualifiedId, backendConfig.domain),
+                        team,
+                        null
                     )
                 }
             } else {
                 val team = httpClient.createTeam(backendConfig)
                 val users = (1..backendConfig.userCount.toInt()).map { httpClient.createUser(it, team, backendConfig.passwordForUsers) }
-                    .plus(team.owner)
                 if (backendConfig.dumpUsers) {
-                    dumpUsers(team, users)
+                    dumpUsers(team, users.plus(team.owner))
                 }
                 users.forEach { setUserHandle(backendConfig, it) }
-                users
+                users.plus(team.owner)
             }
         }
     }
@@ -160,9 +165,10 @@ private suspend fun HttpClient.createTeam(backendConfig: BackendConfig): Team {
         ownerEmail = email,
         ownerPassword = backendConfig.passwordForUsers,
         ownerId = UserId(userId, backend.domain),
+        ownerOldCode = null,
         client = this
     )
-    login(team.owner.email, team.owner.password)
+    login(team.owner.email, team.owner.password, team.owner.request2FA())
 
     put("self/handle") { setBody(mapOf("handle" to ownerId).toJsonObject()) }
     logger.i("Owner $email (id $userId) of team ${backendConfig.teamName} (id: $teamId) in backend ${backendConfig.domain}")
@@ -172,9 +178,11 @@ private suspend fun HttpClient.createTeam(backendConfig: BackendConfig): Team {
 private suspend fun setUserHandle(backendConfig: BackendConfig, userData: UserData) {
     lateinit var token: BearerTokens
     val httpClient = httpClient(backendConfig) { token }
-    httpClient.login(userData.email, userData.password) { accessToken ->
+    val secondFactor = if (backendConfig.secondFactorAuth) httpClient.request2FA(userData.email, userData.userId) else null
+    httpClient.login(userData.email, userData.password, secondFactor) { accessToken ->
         token = BearerTokens(accessToken, "")
     }
+    userData.oldCode = secondFactor
     httpClient.put("self/handle") { setBody(mapOf("handle" to "monkey-${userData.userId.value}").toJsonObject()) }
 }
 
@@ -192,20 +200,21 @@ private suspend fun HttpClient.createUser(i: Int, team: Team, userPassword: Stri
     }.body<JsonObject>()
     val userId = response["id"]?.jsonPrimitive?.content ?: error("Could not register user in team")
     logger.d("Created user $email (id $userId) in team ${team.name}")
-    return UserData(email, userPassword, UserId(userId, team.backend.domain), team)
+    return UserData(email, userPassword, UserId(userId, team.backend.domain), team, null)
 }
 
 private suspend fun HttpClient.login(
     email: String,
     password: String,
+    secondFactor: String?,
     tokenProvider: (accessToken: String) -> Unit = { accessToken ->
         token = BearerTokens(accessToken, "")
     }
 ) {
     val response = post("login") {
         setBody(
             mapOf(
-                "email" to email, "password" to password, "label" to ""
+                "email" to email, "password" to password, "label" to "", "verification_code" to secondFactor
             ).toJsonObject()
         )
     }.body<JsonObject>()
@@ -214,13 +223,24 @@ private suspend fun HttpClient.login(
 }
 
 internal suspend fun HttpClient.teamParticipants(team: Team): List<UserId> {
-    this.login(team.owner.email, team.owner.password)
+    this.login(team.owner.email, team.owner.password, team.owner.request2FA())
     val members = get("teams/${team.id}/members").body<JsonObject>()
     return members["members"]!!.jsonArray.map { member ->
         UserId(member.jsonObject["user"]!!.jsonPrimitive.content, team.backend.domain)
     }
 }
 
+suspend fun HttpClient.request2FA(email: String, userId: UserId): String {
+    this.post("verification-code/send") {
+        setBody(
+            mapOf(
+                "action" to "login", "email" to email
+            ).toJsonObject()
+        )
+    }
+    return this.get("i/users/${userId.value}/verification-code/login").body<JsonPrimitive>().content
+}
+
 private suspend fun HttpClient.invite(teamId: String, email: String, name: String): String {
     val invitationId = post("teams/$teamId/invitations") {
         setBody(
@@ -234,7 +254,7 @@ private suspend fun HttpClient.invite(teamId: String, email: String, name: Strin
 }
 
 internal fun httpClient(backendConfig: BackendConfig, tokenProvider: () -> BearerTokens? = { token }) = HttpClient(OkHttp.create()) {
-    val excludedPaths = listOf("register", "login", "activate")
+    val excludedPaths = listOf("verification-code", "register", "login", "activate")
     defaultRequest {
         url(backendConfig.api)
         contentType(ContentType.Application.Json)

monkeys/src/main/kotlin/com/wire/kalium/monkeys/model/UserData.kt

@@ -19,12 +19,15 @@ package com.wire.kalium.monkeys.model
 
 import com.wire.kalium.logic.data.user.UserId
 import io.ktor.client.HttpClient
+import io.ktor.client.plugins.ClientRequestException
+import io.ktor.http.HttpStatusCode
 
 data class UserData(
     val email: String,
     val password: String,
     val userId: UserId,
     val team: Team,
+    var oldCode: String?
 ) {
     fun backendConfig() = BackendConfig(
         this.team.backend.api,
@@ -37,15 +40,21 @@ data class UserData(
         this.password,
         this.team.backend.domain,
         this.team.name,
-        "",
-        "",
+        this.team.backend.authUser,
+        this.team.backend.authPassword,
         1u,
+        this.team.backend.secondFactorAuthEnabled,
         presetTeam = TeamConfig(
             this.team.id,
             UserAccount(this.team.owner.email, this.team.owner.userId.value),
             users = listOf(UserAccount(this.email, this.userId.value))
         )
     )
+
+    suspend fun request2FA(): String? {
+        this.oldCode = if (this.team.backend.secondFactorAuthEnabled) this.team.request2FA(this.email, this.userId, this.oldCode) else null
+        return this.oldCode
+    }
 }
 
 @Suppress("LongParameterList")
@@ -56,16 +65,27 @@ class Team(
     ownerEmail: String,
     ownerPassword: String,
     ownerId: UserId,
+    ownerOldCode: String?,
     private val client: HttpClient
 ) {
     val owner: UserData
 
     init {
-        this.owner = UserData(ownerEmail, ownerPassword, ownerId, this)
+        this.owner = UserData(ownerEmail, ownerPassword, ownerId, this, ownerOldCode)
     }
 
     suspend fun usersFromTeam(): List<UserId> = this.client.teamParticipants(this)
 
+    internal suspend fun request2FA(email: String, userId: UserId, oldCode: String?) = try {
+        this.client.request2FA(email, userId)
+    } catch (e: ClientRequestException) {
+        if (e.response.status == HttpStatusCode.TooManyRequests) {
+            oldCode
+        } else {
+            throw e
+        }
+    }
+
     override fun equals(other: Any?): Boolean {
         return other != null && other is Team && other.id == this.id
     }
@@ -84,10 +104,13 @@ data class Backend(
     val website: String,
     val title: String,
     val domain: String,
+    val secondFactorAuthEnabled: Boolean,
+    val authUser: String,
+    val authPassword: String,
 ) {
     companion object {
         fun fromConfig(config: BackendConfig): Backend = with(config) {
-            Backend(api, accounts, webSocket, blackList, teams, website, title, domain)
+            Backend(api, accounts, webSocket, blackList, teams, website, title, domain, secondFactorAuth, authUser, authPassword)
         }
     }
 }

monkeys/src/main/kotlin/com/wire/kalium/monkeys/server/MonkeyServer.kt

@@ -49,6 +49,7 @@ class MonkeyServer : CliktCommand() {
     private val port by option("-p", "--port", help = "Port to bind the http server").int().default(DEFAULT_PORT)
     private val logLevel by option("-l", "--log-level", help = "log level").enum<KaliumLogLevel>().default(KaliumLogLevel.INFO)
     private val logOutputFile by option("-o", "--log-file", help = "output file for logs")
+    private val oldCode by option("-c", "--code", help = "Current 2FA code to use until a new one can be generated")
     private val fileLogger: LogWriter by lazy { fileLogger(logOutputFile ?: "kalium.log") }
 
     @OptIn(ExperimentalSerializationApi::class)
@@ -72,11 +73,11 @@ class MonkeyServer : CliktCommand() {
         } else {
             CoreLogger.init(KaliumLogger.Config(logLevel))
         }
-        backendConfig?.let { initMonkey(it) }
+        backendConfig?.let { initMonkey(it, oldCode) }
         embeddedServer(Netty, port = port, host = "0.0.0.0", module = {
             configureMonitoring()
             configureAdministration()
-            configureRoutes(coreLogic)
+            configureRoutes(coreLogic, oldCode)
         }).start(wait = true)
     }
 }

monkeys/src/main/kotlin/com/wire/kalium/monkeys/server/routes/Monkey.kt

@@ -29,7 +29,7 @@ import io.ktor.server.routing.routing
 
 private lateinit var monkey: Monkey
 
-fun initMonkey(backendConfig: BackendConfig) {
+fun initMonkey(backendConfig: BackendConfig, oldCode: String?) {
     val presetTeam = backendConfig.presetTeam ?: error("Preset team must contain exact one user")
     val httpClient = httpClient(backendConfig)
     val backend = Backend.fromConfig(backendConfig)
@@ -40,26 +40,27 @@ fun initMonkey(backendConfig: BackendConfig) {
         presetTeam.owner.email,
         backendConfig.passwordForUsers,
         UserId(presetTeam.owner.unqualifiedId, backendConfig.domain),
+        null,
         httpClient
     )
     val userData = presetTeam.users.map { user ->
         UserData(
-            user.email, backendConfig.passwordForUsers, UserId(user.unqualifiedId, backendConfig.domain), team
+            user.email, backendConfig.passwordForUsers, UserId(user.unqualifiedId, backendConfig.domain), team, oldCode
         )
     }.single()
     // currently the monkey id is not necessary in the server since the coordinator will be the one handling events for the replayer
     monkey = Monkey.internal(userData, MonkeyId.dummy())
 }
 
 @Suppress("LongMethod")
-fun Application.configureRoutes(core: CoreLogic) {
+fun Application.configureRoutes(core: CoreLogic, oldCode: String?) {
     install(ContentNegotiation) {
         json()
     }
     routing {
         post("/$SET_MONKEY") {
             val backendConfig = call.receive<BackendConfig>()
-            initMonkey(backendConfig)
+            initMonkey(backendConfig, oldCode)
             call.respond(HttpStatusCode.OK)
         }
         get("/$IS_SESSION_ACTIVE") {