TheHive-Project#264 OAuth2 refactoring

Mr-P1nqu3 · Aug 12, 2020 · b4836af · b4836af
1 parent 276f79b
commit b4836af
Show file tree

Hide file tree

Showing 8 changed files with 210 additions and 140 deletions.
diff --git a/app/org/thp/cortex/controllers/AuthenticationCtrl.scala b/app/org/thp/cortex/controllers/AuthenticationCtrl.scala
@@ -1,21 +1,21 @@
 package org.thp.cortex.controllers
 
-import scala.concurrent.{ExecutionContext, Future}
-
-import play.api.mvc._
-
 import javax.inject.{Inject, Singleton}
-import org.thp.cortex.models.UserStatus
-import org.thp.cortex.services.UserSrv
-
 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, OAuth2Redirect, Timed}
+import org.elastic4play.{AuthorizationError, MissingAttributeError, Timed}
+import org.thp.cortex.models.UserStatus
+import org.thp.cortex.services.UserSrv
+import play.api.Configuration
+import play.api.mvc._
+
+import scala.concurrent.{ExecutionContext, Future}
 
 @Singleton
 class AuthenticationCtrl @Inject()(
+    configuration: Configuration,
     authSrv: AuthSrv,
     userSrv: UserSrv,
     authenticated: Authenticated,
@@ -44,24 +44,23 @@ class AuthenticationCtrl @Inject()(
     dbIndex.getIndexStatus.flatMap {
       case false ⇒ Future.successful(Results.Status(520))
       case _ ⇒
-        (for {
-          authContext ← authSrv.authenticate()
-          user        ← userSrv.get(authContext.userId)
-        } yield {
-          if (user.status() == UserStatus.Ok)
-            authenticated.setSessingUser(Ok, authContext)
-          else
-            throw AuthorizationError("Your account is locked")
-        }) recover {
-          // A bit of a hack with the status code, so that Angular doesn't reject the origin
-          case OAuth2Redirect(redirectUrl, qp) ⇒ Redirect(redirectUrl, qp, status = OK)
-          case e                               ⇒ throw e
-        }
+        authSrv
+          .authenticate()
+          .flatMap {
+            case Right(authContext) ⇒
+              userSrv.get(authContext.userId).map { user ⇒
+                if (user.status() == UserStatus.Ok)
+                  authenticated.setSessingUser(Redirect(configuration.get[String]("play.http.context").stripSuffix("/") + "/index.html"), authContext)
+                else
+                  throw AuthorizationError("Your account is locked")
+              }
+            case Left(result) ⇒ Future.successful(result)
+          }
     }
   }
 
   @Timed
-  def logout = Action {
+  def logout: Action[AnyContent] = Action {
     Ok.withNewSession
   }
 }
diff --git a/app/org/thp/cortex/controllers/JobCtrl.scala b/app/org/thp/cortex/controllers/JobCtrl.scala
@@ -2,6 +2,7 @@ package org.thp.cortex.controllers
 
 import scala.concurrent.duration.{Duration, FiniteDuration}
 import scala.concurrent.{ExecutionContext, Future}
+import scala.concurrent.duration.DurationInt
 
 import play.api.http.Status
 import play.api.libs.json.{JsObject, JsString, JsValue, Json}
@@ -146,7 +147,7 @@ class JobCtrl @Inject()(
       .flatMap {
         case job if job.status() == JobStatus.InProgress || job.status() == JobStatus.Waiting ⇒
           val duration                  = Duration(atMost).asInstanceOf[FiniteDuration]
-          implicit val timeout: Timeout = Timeout(duration)
+          implicit val timeout: Timeout = Timeout(duration + 1.second)
           (auditActor ? Register(jobId, duration))
             .mapTo[JobEnded]
             .map(_ ⇒ ())

diff --git a/app/org/thp/cortex/services/OAuth2Srv.scala b/app/org/thp/cortex/services/OAuth2Srv.scala
@@ -1,19 +1,18 @@
 package org.thp.cortex.services
-
-import scala.concurrent.{ExecutionContext, Future}
-
-import play.api.http.Status
-import play.api.libs.json.{JsObject, JsValue}
-import play.api.libs.ws.WSClient
-import play.api.mvc.RequestHeader
-import play.api.{Configuration, Logger}
+import java.util.UUID
 
 import akka.stream.Materializer
 import javax.inject.{Inject, Singleton}
+import org.elastic4play.services.{AuthContext, AuthSrv}
+import org.elastic4play.{AuthenticationError, BadRequestError, NotFoundError}
 import org.thp.cortex.services.mappers.UserMapper
+import play.api.libs.json.JsObject
+import play.api.libs.ws.WSClient
+import play.api.mvc.{RequestHeader, Result, Results}
+import play.api.{Configuration, Logger}
 
-import org.elastic4play.services.{AuthContext, AuthSrv}
-import org.elastic4play.{AuthenticationError, AuthorizationError, OAuth2Redirect}
+import scala.concurrent.{ExecutionContext, Future}
+import scala.util.{Failure, Success}
 
 case class OAuth2Config(
     clientId: String,
@@ -25,24 +24,41 @@ case class OAuth2Config(
     tokenUrl: String,
     userUrl: String,
     scope: String,
+    authorizationHeader: String,
+    autoupdate: Boolean,
     autocreate: Boolean
 )
 
 object OAuth2Config {
 
   def apply(configuration: Configuration): Option[OAuth2Config] =
     for {
-      clientId         ← configuration.getOptional[String]("auth.oauth2.clientId")
-      clientSecret     ← configuration.getOptional[String]("auth.oauth2.clientSecret")
-      redirectUri      ← configuration.getOptional[String]("auth.oauth2.redirectUri")
-      responseType     ← configuration.getOptional[String]("auth.oauth2.responseType")
-      grantType        ← configuration.getOptional[String]("auth.oauth2.grantType")
+      clientId     ← configuration.getOptional[String]("auth.oauth2.clientId")
+      clientSecret ← configuration.getOptional[String]("auth.oauth2.clientSecret")
+      redirectUri  ← configuration.getOptional[String]("auth.oauth2.redirectUri")
+      responseType ← configuration.getOptional[String]("auth.oauth2.responseType")
+      grantType = configuration.getOptional[String]("auth.oauth2.grantType").getOrElse("authorization_code")
       authorizationUrl ← configuration.getOptional[String]("auth.oauth2.authorizationUrl")
-      userUrl          ← configuration.getOptional[String]("auth.oauth2.userUrl")
       tokenUrl         ← configuration.getOptional[String]("auth.oauth2.tokenUrl")
+      userUrl          ← configuration.getOptional[String]("auth.oauth2.userUrl")
       scope            ← configuration.getOptional[String]("auth.oauth2.scope")
-      autocreate = configuration.getOptional[Boolean]("auth.sso.autocreate").getOrElse(false)
-    } yield OAuth2Config(clientId, clientSecret, redirectUri, responseType, grantType, authorizationUrl, tokenUrl, userUrl, scope, autocreate)
+      authorizationHeader = configuration.getOptional[String]("auth.oauth2.authorizationHeader").getOrElse("Bearer")
+      autocreate          = configuration.getOptional[Boolean]("auth.sso.autocreate").getOrElse(false)
+      autoupdate          = configuration.getOptional[Boolean]("auth.sso.autoupdate").getOrElse(false)
+    } yield OAuth2Config(
+      clientId,
+      clientSecret,
+      redirectUri,
+      responseType,
+      grantType,
+      authorizationUrl,
+      tokenUrl,
+      userUrl,
+      scope,
+      authorizationHeader,
+      autocreate,
+      autoupdate
+    )
 }
 
 @Singleton
@@ -66,92 +82,146 @@ class OAuth2Srv(
   private def withOAuth2Config[A](body: OAuth2Config ⇒ Future[A]): Future[A] =
     oauth2Config.fold[Future[A]](Future.failed(AuthenticationError("OAuth2 not configured properly")))(body)
 
-  override def authenticate()(implicit request: RequestHeader): Future[AuthContext] =
-    withOAuth2Config { cfg ⇒
-      request
-        .queryString
-        .get(Oauth2TokenQueryString)
-        .flatMap(_.headOption)
-        .fold(createOauth2Redirect(cfg.clientId)) { code ⇒
-          getAuthTokenAndAuthenticate(cfg.clientId, code)
+  override def authenticate()(implicit request: RequestHeader): Future[Either[Result, AuthContext]] =
+    withOAuth2Config { oauth2Config ⇒
+      if (!isSecuredAuthCode(request)) {
+        logger.debug("Code or state is not provided, redirect to authorizationUrl")
+        Future.successful(Left(authRedirect(oauth2Config)))
+      } else {
+        (for {
+          token       ← getToken(oauth2Config, request)
+          userData    ← getUserData(oauth2Config, token)
+          authContext ← authenticate(oauth2Config, request, userData)
+        } yield Right(authContext)).recoverWith {
+          case error ⇒ Future.failed(AuthenticationError(s"OAuth2 authentication failure: ${error.getMessage}"))
         }
+      }
     }
 
-  private def getAuthTokenAndAuthenticate(clientId: String, code: String)(implicit request: RequestHeader): Future[AuthContext] = {
-    logger.debug("Getting user token with the code from the response!")
-    withOAuth2Config { cfg ⇒
-      val acceptHeader = "Accept" → cfg.responseType
-      ws.url(cfg.tokenUrl)
-        .addHttpHeaders(acceptHeader)
-        .post(
-          Map(
-            "code"          → code,
-            "grant_type"    → cfg.grantType,
-            "client_secret" → cfg.clientSecret,
-            "redirect_uri"  → cfg.redirectUri,
-            "client_id"     → clientId
-          )
-        )
-        .recoverWith {
-          case error ⇒
-            logger.error(s"Token verification failure", error)
-            Future.failed(AuthenticationError("Token verification failure"))
-        }
-        .flatMap { r ⇒
-          r.status match {
-            case Status.OK ⇒
-              val accessToken = (r.json \ "access_token").asOpt[String].getOrElse("")
-              val authHeader  = "Authorization" → s"bearer $accessToken"
-              ws.url(cfg.userUrl)
-                .addHttpHeaders(authHeader)
-                .get()
-                .flatMap { userResponse ⇒
-                  if (userResponse.status != Status.OK) {
-                    Future.failed(AuthenticationError(s"unexpected response from server: ${userResponse.status} ${userResponse.body}"))
-                  } else {
-                    val response = userResponse.json.asInstanceOf[JsObject]
-                    getOrCreateUser(response, authHeader)
-                  }
-                }
-            case _ ⇒
-              logger.error(s"unexpected response from server: ${r.status} ${r.body}")
-              Future.failed(AuthenticationError("unexpected response from server"))
-          }
-        }
-    }
+  private def isSecuredAuthCode(request: RequestHeader): Boolean =
+    request.queryString.contains("code") && request.queryString.contains("state")
+
+  /**
+    * Filter checking whether we initiate the OAuth2 process
+    * and redirecting to OAuth2 server if necessary
+    * @return
+    */
+  private def authRedirect(oauth2Config: OAuth2Config): Result = {
+    val state = UUID.randomUUID().toString
+    val queryStringParams = Map[String, Seq[String]](
+      "scope"         → Seq(oauth2Config.scope),
+      "response_type" → Seq(oauth2Config.responseType),
+      "redirect_uri"  → Seq(oauth2Config.redirectUri),
+      "client_id"     → Seq(oauth2Config.clientId),
+      "state"         → Seq(state)
+    )
+
+    logger.debug(s"Redirecting to ${oauth2Config.redirectUri} with $queryStringParams and state $state")
+    Results
+      .Redirect(oauth2Config.authorizationUrl, queryStringParams, status = 302)
+      .withSession("state" → state)
   }
 
-  private def getOrCreateUser(response: JsValue, authHeader: (String, String))(implicit request: RequestHeader): Future[AuthContext] =
-    withOAuth2Config { cfg ⇒
-      ssoMapper.getUserFields(response, Some(authHeader)).flatMap { userFields ⇒
-        val userId = userFields.getString("login").getOrElse("")
-        userSrv
-          .get(userId)
-          .flatMap(user ⇒ {
-            userSrv.getFromUser(request, user, name)
-          })
-          .recoverWith {
-            case authErr: AuthorizationError ⇒ Future.failed(authErr)
-            case _ if cfg.autocreate ⇒
-              userSrv.inInitAuthContext { implicit authContext ⇒
-                userSrv
-                  .create(userFields)
-                  .flatMap(user ⇒ {
-                    userSrv.getFromUser(request, user, name)
-                  })
-              }
-          }
+  /**
+    * Enriching the initial request with OAuth2 token gotten
+    * from OAuth2 code
+    * @return
+    */
+  private def getToken[A](oauth2Config: OAuth2Config, request: RequestHeader): Future[String] = {
+    val token =
+      for {
+        state   ← request.session.get("state")
+        stateQs ← request.queryString.get("state").flatMap(_.headOption)
+        if state == stateQs
+      } yield request.queryString.get("code").flatMap(_.headOption) match {
+        case Some(code) ⇒
+          logger.debug(s"Attempting to retrieve OAuth2 token from ${oauth2Config.tokenUrl} with code $code")
+          getAuthTokenFromCode(oauth2Config, code, state)
+            .map { t ⇒
+              logger.trace(s"Got token $t")
+              t
+            }
+        case None ⇒
+          Future.failed(AuthenticationError(s"OAuth2 server code missing ${request.queryString.get("error")}"))
       }
-    }
+    token.getOrElse(Future.failed(BadRequestError("OAuth2 states mismatch")))
+  }
 
-  private def createOauth2Redirect(clientId: String): Future[AuthContext] =
-    withOAuth2Config { cfg ⇒
-      val queryStringParams = Map[String, Seq[String]](
-        "scope"         → Seq(cfg.scope),
-        "response_type" → Seq(cfg.responseType),
-        "redirect_uri"  → Seq(cfg.redirectUri),
-        "client_id"     → Seq(clientId)
+  /**
+    * Querying the OAuth2 server for a token
+    * @param code the previously obtained code
+    * @return
+    */
+  private def getAuthTokenFromCode(oauth2Config: OAuth2Config, code: String, state: String): Future[String] = {
+    logger.trace(s"""
+                    |Request to ${oauth2Config.tokenUrl} with
+                    |  code:          $code
+                    |  grant_type:    ${oauth2Config.grantType}
+                    |  client_secret: ${oauth2Config.clientSecret}
+                    |  redirect_uri:  ${oauth2Config.redirectUri}
+                    |  client_id:     ${oauth2Config.clientId}
+                    |  state:         $state
+                    |""".stripMargin)
+    ws.url(oauth2Config.tokenUrl)
+      .withHttpHeaders("Accept" → "application/json")
+      .post(
+        Map(
+          "code"          → code,
+          "grant_type"    → oauth2Config.grantType,
+          "client_secret" → oauth2Config.clientSecret,
+          "redirect_uri"  → oauth2Config.redirectUri,
+          "client_id"     → oauth2Config.clientId,
+          "state"         → state
+        )
       )
-      Future.failed(OAuth2Redirect(cfg.authorizationUrl, queryStringParams))
-    }
+      .transform {
+        case Success(r) if r.status == 200 ⇒ Success((r.json \ "access_token").asOpt[String].getOrElse(""))
+        case Failure(error)                ⇒ Failure(AuthenticationError(s"OAuth2 token verification failure ${error.getMessage}"))
+        case Success(r)                    ⇒ Failure(AuthenticationError(s"OAuth2/token unexpected response from server (${r.status} ${r.statusText})"))
+      }
+  }
+
+  /**
+    * Client query for user data with OAuth2 token
+    * @param token the token
+    * @return
+    */
+  private def getUserData(oauth2Config: OAuth2Config, token: String): Future[JsObject] = {
+    logger.trace(s"Request to ${oauth2Config.userUrl} with authorization header: ${oauth2Config.authorizationHeader} $token")
+    ws.url(oauth2Config.userUrl)
+      .addHttpHeaders("Authorization" → s"${oauth2Config.authorizationHeader} $token")
+      .get()
+      .transform {
+        case Success(r) if r.status == 200 ⇒ Success(r.json.as[JsObject])
+        case Failure(error)                ⇒ Failure(AuthenticationError(s"OAuth2 user data fetch failure ${error.getMessage}"))
+        case Success(r)                    ⇒ Failure(AuthenticationError(s"OAuth2/userinfo unexpected response from server (${r.status} ${r.statusText})"))
+      }
+  }
+
+  private def authenticate(oauth2Config: OAuth2Config, request: RequestHeader, userData: JsObject): Future[AuthContext] =
+    for {
+      userFields ← ssoMapper.getUserFields(userData)
+      login      ← userFields.getString("login").fold(Future.failed[String](AuthenticationError("")))(Future.successful)
+      user ← userSrv
+        .get(login)
+        .flatMap {
+          case u if oauth2Config.autoupdate ⇒
+            logger.debug(s"Updating OAuth/OIDC user")
+            userSrv.inInitAuthContext { implicit authContext ⇒
+              // Only update name and roles, not login (can't change it)
+              userSrv
+                .update(u, userFields.unset("login"))
+
+            }
+          case u ⇒ Future.successful(u)
+        }
+        .recoverWith {
+          case _: NotFoundError if oauth2Config.autocreate ⇒
+            logger.debug(s"Creating OAuth/OIDC user")
+            userSrv.inInitAuthContext { implicit authContext ⇒
+              userSrv.create(userFields.set("login", userFields.getString("login").get.toLowerCase))
+            }
+        }
+      authContext ← userSrv.getFromUser(request, user, name)
+    } yield authContext
 }