feat: optional TOTP multi-factor authentication

app/org/thp/cortex/controllers/AuthenticationCtrl.scala

@@ -3,11 +3,10 @@ package org.thp.cortex.controllers
 import javax.inject.{Inject, Singleton}
 import org.elastic4play.controllers.{Authenticated, Fields, FieldsBodyParser, Renderer}
 import org.elastic4play.database.DBIndex
-import org.elastic4play.services.AuthSrv
 import org.elastic4play.services.JsonFormat.authContextWrites
 import org.elastic4play.{AuthorizationError, MissingAttributeError, Timed}
 import org.thp.cortex.models.UserStatus
-import org.thp.cortex.services.UserSrv
+import org.thp.cortex.services.{CortexAuthSrv, UserSrv}
 import play.api.Configuration
 import play.api.mvc._
 
@@ -16,7 +15,8 @@ import scala.concurrent.{ExecutionContext, Future}
 @Singleton
 class AuthenticationCtrl @Inject() (
     configuration: Configuration,
-    authSrv: AuthSrv,
+    // CortexAuthSrv rather than the AuthSrv trait: the second-factor-aware authenticate overload lives there
+    authSrv: CortexAuthSrv,
     userSrv: UserSrv,
     authenticated: Authenticated,
     dbIndex: DBIndex,
@@ -32,9 +32,12 @@ class AuthenticationCtrl @Inject() (
       case false => Future.successful(Results.Status(520))
       case _ =>
         for {
-          user        <- request.body.getString("user").fold[Future[String]](Future.failed(MissingAttributeError("user")))(Future.successful)
-          password    <- request.body.getString("password").fold[Future[String]](Future.failed(MissingAttributeError("password")))(Future.successful)
-          authContext <- authSrv.authenticate(user, password)
+          user     <- request.body.getString("user").fold[Future[String]](Future.failed(MissingAttributeError("user")))(Future.successful)
+          password <- request.body.getString("password").fold[Future[String]](Future.failed(MissingAttributeError("password")))(Future.successful)
+          // Optional second factor. Absent on the first attempt; the client retries with it after a
+          // MultiFactorCodeRequired error. Also accepts a single-use backup code.
+          code = request.body.getString("code")
+          authContext <- authSrv.authenticate(user, password, code)
         } yield authenticated.setSessingUser(renderer.toOutput(OK, authContext), authContext)
     }
   }

app/org/thp/cortex/controllers/UserCtrl.scala

@@ -6,7 +6,7 @@ import org.elastic4play.models.JsonFormat.baseModelEntityWrites
 import org.elastic4play.services.JsonFormat.queryReads
 import org.elastic4play.services.{AuthContext, AuthSrv, QueryDSL, QueryDef}
 import org.thp.cortex.models.{OrganizationStatus, Roles}
-import org.thp.cortex.services.{OrganizationSrv, UserSrv}
+import org.thp.cortex.services.{OrganizationSrv, TOTPSrv, UserSrv}
 import play.api.Logger
 import play.api.http.Status
 import play.api.libs.json.{JsObject, Json}
@@ -20,6 +20,7 @@ import scala.util.Try
 class UserCtrl @Inject() (
     userSrv: UserSrv,
     authSrv: AuthSrv,
+    totpSrv: TOTPSrv,
     organizationSrv: OrganizationSrv,
     authenticated: Authenticated,
     renderer: Renderer,
@@ -112,6 +113,8 @@ class UserCtrl @Inject() (
       if (request.body.contains("password"))
         Future.failed(AuthorizationError("You must use dedicated API (setPassword, changePassword) to update password"))
       else if (request.body.contains("key")) Future.failed(AuthorizationError("You must use dedicated API (renewKey, removeKey) to update key"))
+      else if (request.body.contains("totpSecret") || request.body.contains("totpScratchCodes"))
+        Future.failed(AuthorizationError("You must use dedicated API (mfa/init, mfa/set, mfa/unset) to update multi-factor authentication"))
       else Future.successful(())
 
     for {
@@ -254,4 +257,68 @@ class UserCtrl @Inject() (
       key      <- authSrv.setKey(userId, keyInput)
     } yield Ok(key)
   }
+
+  /** Session key holding the secret between `mfaInit` and `mfaSet`.
+    *
+    * Keeping the pending secret in the (signed) session rather than server-side memory keeps enrolment working across a horizontally scaled Cortex
+    * without sticky sessions, and stops the client from choosing its own secret. The value is readable by the user's own browser, which is fine —
+    * it is the very thing they are about to scan.
+    */
+  private val pendingMfaSecret = "mfaSecret"
+
+  @Timed
+  def mfaInit(userId: String): Action[AnyContent] = authenticated().async { implicit request =>
+    if (userId != request.authContext.userId)
+      Future.failed(AuthorizationError("You can't enable multi-factor authentication for another user"))
+    else if (!totpSrv.enabled)
+      Future.failed(BadRequestError("Multi-factor authentication is disabled on this instance"))
+    else {
+      val secret = totpSrv.generateSecret()
+      Future.successful(
+        renderer
+          .toOutput(
+            OK,
+            Json.obj(
+              "secret" -> secret,                             // for manual entry
+              "uri"    -> totpSrv.otpAuthUri(userId, secret), // for a click-through on mobile
+              "qrCode" -> totpSrv.qrCode(userId, secret)      // SVG data URI, for an <img>
+            )
+          )
+          .addingToSession(pendingMfaSecret -> secret)
+      )
+    }
+  }
+
+  @Timed
+  def mfaSet(userId: String): Action[Fields] = authenticated().async(fieldsBodyParser) { implicit request =>
+    if (userId != request.authContext.userId)
+      Future.failed(AuthorizationError("You can't enable multi-factor authentication for another user"))
+    else
+      for {
+        secret <- request
+          .session
+          .get(pendingMfaSecret)
+          .fold(Future.failed[String](BadRequestError("No multi-factor enrolment in progress, start it again")))(Future.successful)
+        code <- request.body.getString("code").fold(Future.failed[String](MissingAttributeError("code")))(Future.successful)
+        _    <- totpSrv.verifyForEnrolment(userId, secret, code)
+        (backupCodes, backupCodeHashes) = totpSrv.generateBackupCodes()
+        _ <- totpSrv.enable(userId, secret, backupCodeHashes)
+        _ = logger.info(s"User $userId enabled multi-factor authentication")
+      } yield renderer
+        .toOutput(OK, Json.obj("backupCodes" -> backupCodes))
+        .removingFromSession(pendingMfaSecret)
+  }
+
+  /** Self-service disable, and the admin reset for a user who lost their authenticator. */
+  @Timed
+  def mfaUnset(userId: String): Action[AnyContent] = authenticated().async { implicit request =>
+    val isAdmin = request.roles.contains(Roles.orgAdmin) || request.roles.contains(Roles.superAdmin)
+    for {
+      _ <- if (userId == request.authContext.userId || isAdmin) Future.successful(())
+      else Future.failed(AuthorizationError("You are not authorized to perform this operation"))
+      _ <- checkUserOrganization(userId)
+      _ <- totpSrv.disable(userId)
+      _ = logger.info(s"Multi-factor authentication of user $userId reset by user ${request.userId}")
+    } yield NoContent
+  }
 }

app/org/thp/cortex/models/Errors.scala

@@ -9,3 +9,9 @@ case class RateLimitExceeded(analyzer: Worker)
     extends CortexError(
       s"Rate limit of ${analyzer.rate().getOrElse("(not set ?!)")} per ${analyzer.rateUnit().getOrElse("(not set ?!)")} reached for the analyzer ${analyzer.name()}. Job cannot be started"
     )
+
+/** The credentials are valid but the user has MFA enabled and didn't supply a code. Clients must retry the login with a "code" field. */
+case class MultiFactorCodeRequired(message: String = "Multi-factor authentication code is required") extends CortexError(message)
+
+/** The credentials are valid but the supplied MFA code is wrong, or too many wrong codes have been submitted. */
+case class MultiFactorCodeInvalid(message: String = "Multi-factor authentication code is invalid") extends CortexError(message)

app/org/thp/cortex/models/User.scala

@@ -1,7 +1,7 @@
 package org.thp.cortex.models
 
 import scala.concurrent.Future
-import play.api.libs.json.{Format, JsArray, JsBoolean, JsObject, JsString}
+import play.api.libs.json.{Format, JsArray, JsBoolean, JsNumber, JsObject, JsString}
 import org.elastic4play.models.JsonFormat.enumFormat
 import org.elastic4play.models.{AttributeDef, BaseEntity, EntityDef, HiveEnumeration, ModelDef, AttributeFormat => F, AttributeOption => O}
 import org.elastic4play.services.{User => EUser}
@@ -23,6 +23,9 @@ trait UserAttributes { _: AttributeDef =>
   val avatar       = optionalAttribute("avatar", F.rawFmt, "Base64 representation of user avatar image", O.unaudited)
   val preferences  = attribute("preferences", F.rawFmt, "User preferences", "{}", O.sensitive, O.unaudited)
   val organization = attribute("organization", F.stringFmt, "User organization")
+  val totpSecret   = optionalAttribute("totpSecret", F.stringFmt, "TOTP shared secret (base32)", O.sensitive, O.unaudited)
+  // Single-use backup codes, consumed one by one when the user can't produce a TOTP code
+  val totpScratchCodes = multiAttribute("totpScratchCodes", F.stringFmt, "Unused MFA backup codes", O.sensitive, O.unaudited)
 }
 
 class UserModel extends ModelDef[UserModel, User]("user", "User", "/user") with UserAttributes with AuditedModel {
@@ -40,7 +43,9 @@ class User(model: UserModel, attributes: JsObject) extends EntityDef[UserModel,
 
   override def toJson: JsObject =
     super.toJson +
-      ("roles"       -> JsArray(roles().map(r => JsString(r.name.toLowerCase())))) +
-      ("hasKey"      -> JsBoolean(key().isDefined)) +
-      ("hasPassword" -> JsBoolean(password().isDefined))
+      ("roles"                -> JsArray(roles().map(r => JsString(r.name.toLowerCase())))) +
+      ("hasKey"               -> JsBoolean(key().isDefined)) +
+      ("hasPassword"          -> JsBoolean(password().isDefined)) +
+      ("hasMFA"               -> JsBoolean(totpSecret().isDefined)) +
+      ("remainingBackupCodes" -> JsNumber(totpScratchCodes().size))
 }

app/org/thp/cortex/models/package.scala

@@ -1,5 +1,22 @@
 package org.thp.cortex
 
 package object models {
+
+  /** Version of the stored document shape. Bumping it makes `POST /api/maintenance/migrate` create a new index from the current models and copy
+    * every document into it, so a bump is also an upgrade step every operator has to run before Cortex serves again.
+    *
+    * Deliberately *not* bumped for the multi-factor attributes (`totpSecret`, `totpScratchCodes`). A full reindex of every job, report and artifact
+    * is a steep price for two fields that are only ever read back from `_source` — nothing queries, sorts or aggregates on them. On an index created
+    * before they existed, Elasticsearch maps them on first write (`text` with a `keyword` sub-field, rather than the declared `keyword`), which
+    * changes nothing for the feature.
+    *
+    * If a future change needs to *query* either field, bump this and add the matching `DatabaseState` case in `Migration.scala`. An `exists` query —
+    * the natural way to ask "who has MFA enabled" — matches under both mapping shapes; only an exact `term` on the value needs the declared
+    * `keyword`.
+    *
+    * Test coverage behind the above: enrolment and login were exercised end to end against a freshly created index. The pre-existing-index case was
+    * only checked at the Elasticsearch level — the dynamic mapping it produces, and a `_source` round-trip — not by running Cortex against an index
+    * created by an earlier release. Upgrading an installed package to this build was not tested.
+    */
   val modelVersion = 6
 }

app/org/thp/cortex/services/CortexAuthSrv.scala

@@ -3,11 +3,12 @@ package org.thp.cortex.services
 import javax.inject.{Inject, Singleton}
 
 import scala.collection.immutable
-import scala.concurrent.ExecutionContext
+import scala.concurrent.{ExecutionContext, Future}
 
 import play.api.{Configuration, Logger}
+import play.api.mvc.RequestHeader
 
-import org.elastic4play.services.AuthSrv
+import org.elastic4play.services.{AuthCapability, AuthContext, AuthSrv}
 import org.elastic4play.services.auth.MultiAuthSrv
 
 object CortexAuthSrv {
@@ -30,13 +31,46 @@ class CortexAuthSrv @Inject() (
     configuration: Configuration,
     authModules: immutable.Set[AuthSrv],
     userSrv: UserSrv,
+    totpSrv: TOTPSrv,
     implicit override val ec: ExecutionContext
 ) extends MultiAuthSrv(
       CortexAuthSrv.getAuthSrv(configuration.getDeprecated[Option[Seq[String]]]("auth.provider", "auth.type").getOrElse(Seq("local")), authModules),
       ec
     ) {
 
-  // Uncomment the following lines if you want to prevent user with key to use password to authenticate
+  // Recomputed from the providers rather than delegated to MultiAuthSrv: `capabilities` is a val, so
+  // it can't be reached through `super`. The front-end reads this from /api/status to decide whether
+  // to offer MFA enrolment in the settings page.
+  override val capabilities: Set[AuthCapability.Type] = {
+    val providerCapabilities = authProviders.flatMap(_.capabilities).toSet
+    if (totpSrv.enabled) providerCapabilities + AuthCapability.mfa else providerCapabilities
+  }
+
+  /** Password authentication with a second factor.
+    *
+    * The MFA gate deliberately sits *outside* `MultiAuthSrv`'s provider fold: inside it, any failure is swallowed and falls through to the next
+    * provider, which would turn a rejected MFA code into a silent retry. Here the password is checked by whichever provider owns it (local, ldap,
+    * ad) and only then is the code verified, so an MFA failure propagates to the client.
+    *
+    * API keys and SSO are intentionally not gated: keys are how TheHive and MISP integrate, and with SSO the identity provider owns MFA.
+    */
+  def authenticate(username: String, password: String, code: Option[String])(implicit request: RequestHeader): Future[AuthContext] =
+    super
+      .authenticate(username, password)
+      .flatMap { authContext =>
+        userSrv
+          .get(authContext.userId)
+          .flatMap(user => totpSrv.checkAtLogin(user, code))
+          .map(_ => authContext)
+      }
+
+  /** Password authentication with no way to supply a code — HTTP Basic auth. An enrolled user is refused rather than silently exempted. */
+  override def authenticate(username: String, password: String)(implicit request: RequestHeader): Future[AuthContext] =
+    authenticate(username, password, None)
+
+  // Uncomment the following lines if you want to prevent user with key to use password to authenticate.
+  // NOTE: authenticate(username, password) is now overridden above, so fold this logic into that
+  // override rather than pasting a second one.
   //  override def authenticate(username: String, password: String)(implicit request: RequestHeader): Future[AuthContext] =
   //    userSrv.get(username)
   //      .transformWith {

app/org/thp/cortex/services/ErrorHandler.scala

@@ -2,7 +2,7 @@ package org.thp.cortex.services
 
 import org.elastic4play.JsonFormat.attributeCheckingExceptionWrites
 import org.elastic4play._
-import org.thp.cortex.models.{JobNotFoundError, RateLimitExceeded, WorkerNotFoundError}
+import org.thp.cortex.models.{JobNotFoundError, MultiFactorCodeInvalid, MultiFactorCodeRequired, RateLimitExceeded, WorkerNotFoundError}
 import play.api.Logger
 import play.api.http.Status.{BAD_REQUEST, FORBIDDEN, NOT_FOUND}
 import play.api.http.{HttpErrorHandler, Status, Writeable}
@@ -33,8 +33,13 @@ class ErrorHandler extends HttpErrorHandler {
       case AuthorizationError(message)  => Status.FORBIDDEN    -> Json.obj("type" -> "AuthorizationError", "message"  -> message)
       case UpdateError(_, message, attributes) =>
         Status.INTERNAL_SERVER_ERROR -> Json.obj("type" -> "UpdateError", "message" -> message, "object" -> attributes)
-      case rle: RateLimitExceeded => Status.TOO_MANY_REQUESTS     -> Json.obj("type" -> "RateLimitExceeded", "message" -> rle.getMessage)
-      case InternalError(message) => Status.INTERNAL_SERVER_ERROR -> Json.obj("type" -> "InternalError", "message"     -> message)
+      case rle: RateLimitExceeded => Status.TOO_MANY_REQUESTS -> Json.obj("type" -> "RateLimitExceeded", "message" -> rle.getMessage)
+      // The password was correct but the second factor is missing or wrong. Both are 401 so that
+      // a client which knows nothing about MFA still treats them as an authentication failure;
+      // the "type" field is what lets an MFA-aware client tell "ask for a code" from "code refused".
+      case mfa: MultiFactorCodeRequired => Status.UNAUTHORIZED          -> Json.obj("type" -> "MultiFactorCodeRequired", "message" -> mfa.getMessage)
+      case mfa: MultiFactorCodeInvalid  => Status.UNAUTHORIZED          -> Json.obj("type" -> "MultiFactorCodeInvalid", "message"  -> mfa.getMessage)
+      case InternalError(message)       => Status.INTERNAL_SERVER_ERROR -> Json.obj("type" -> "InternalError", "message"           -> message)
       case nfe: NumberFormatException =>
         Status.BAD_REQUEST -> Json.obj("type" -> "NumberFormatException", "message" -> ("Invalid format " + nfe.getMessage))
       case NotFoundError(message)        => Status.NOT_FOUND   -> Json.obj("type" -> "NotFoundError", "message" -> message)