From a8611b5eaa9aa72a0de887fdbd729ac74650fb51 Mon Sep 17 00:00:00 2001
From: Robin Riclet <robinriclet@gmail.com>
Date: Tue, 23 Mar 2021 15:16:09 +0100
Subject: [PATCH] #1872 Added quota for user creation

---
 .../thp/thehive/migration/th4/Output.scala    |  2 +-
 .../org/thp/thehive/services/UserSrv.scala    | 51 ++++++++++++-------
 2 files changed, 35 insertions(+), 18 deletions(-)

diff --git a/migration/src/main/scala/org/thp/thehive/migration/th4/Output.scala b/migration/src/main/scala/org/thp/thehive/migration/th4/Output.scala
index c1cc76b06b..a527f8bd62 100644
--- a/migration/src/main/scala/org/thp/thehive/migration/th4/Output.scala
+++ b/migration/src/main/scala/org/thp/thehive/migration/th4/Output.scala
@@ -326,7 +326,7 @@ class Output @Inject() (
   override def createUser(inputUser: InputUser): Try[IdMapping] =
     authTransaction(inputUser.metaData.createdBy) { implicit graph => implicit authContext =>
       logger.debug(s"Create user ${inputUser.user.login}")
-      userSrv.checkUser(inputUser.user).flatMap(userSrv.createEntity).map { createdUser =>
+      userSrv.checkUserLogin(inputUser.user).flatMap(userSrv.createEntity).map { createdUser =>
         updateMetaData(createdUser, inputUser.metaData)
         inputUser
           .avatar
diff --git a/thehive/app/org/thp/thehive/services/UserSrv.scala b/thehive/app/org/thp/thehive/services/UserSrv.scala
index dcf8255262..108d701a73 100644
--- a/thehive/app/org/thp/thehive/services/UserSrv.scala
+++ b/thehive/app/org/thp/thehive/services/UserSrv.scala
@@ -23,7 +23,7 @@ import play.api.libs.json.{JsObject, Json}
 
 import java.util.regex.Pattern
 import java.util.{List => JList, Map => JMap}
-import javax.inject.{Inject, Named, Singleton}
+import javax.inject.{Inject, Named, Provider, Singleton}
 import scala.util.{Failure, Success, Try}
 
 @Singleton
@@ -32,14 +32,33 @@ class UserSrv @Inject() (
     roleSrv: RoleSrv,
     auditSrv: AuditSrv,
     attachmentSrv: AttachmentSrv,
+    organisationSrvProvider: Provider[OrganisationSrv],
     @Named("integrity-check-actor") integrityCheckActor: ActorRef
 ) extends VertexSrv[User] {
+  lazy val organisationSrv: OrganisationSrv = organisationSrvProvider.get
+
   val defaultUserDomain: Option[String] = configuration.getOptional[String]("auth.defaultUserDomain")
   val fullUserNameRegex: Pattern        = "[\\p{Graph}&&[^@.]](?:[\\p{Graph}&&[^@]]*)*@\\p{Alnum}+(?:[\\p{Alnum}-.])*".r.pattern
 
   val userAttachmentSrv = new EdgeSrv[UserAttachment, User, Attachment]
 
-  def checkUser(user: User): Try[User] = {
+  def addOrCreateUser(user: User, avatar: Option[FFile], organisation: Organisation with Entity, profile: Profile with Entity)(implicit
+      graph: Graph,
+      authContext: AuthContext
+  ): Try[RichUser] =
+    getByName(user.login)
+      .getOrFail("User")
+      .orElse {
+        for {
+          validUser   <- checkUserLogin(user)
+          _           <- checkUserQuota(organisation)
+          createdUser <- createEntity(validUser)
+          _           <- avatar.map(setAvatar(createdUser, _)).flip
+        } yield createdUser
+      }
+      .flatMap(addUserToOrganisation(_, organisation, profile))
+
+  def checkUserLogin(user: User): Try[User] = {
     val login =
       if (!user.login.contains('@') && defaultUserDomain.isDefined) s"${user.login}@${defaultUserDomain.get}".toLowerCase
       else user.login.toLowerCase
@@ -48,6 +67,19 @@ class UserSrv @Inject() (
     else Failure(BadRequestError(s"User login is invalid, it must be an email address (found: ${user.login})"))
   }
 
+  def checkUserQuota(organisation: Organisation with Entity)(implicit
+      graph: Graph,
+      authContext: AuthContext
+  ): Try[Unit] = {
+    val userQuota = configuration.getOptional[Long]("quota.organisation.user.count")
+    val userCount = organisationSrv.get(organisation).users.getCount
+
+    userQuota.fold[Try[Unit]](Success(()))(quota =>
+      if (userCount < quota) Success(())
+      else Failure(BadRequestError(s"User quota is reached, this organisation cannot have more users"))
+    )
+  }
+
   // TODO return Try[Unit]
   def addUserToOrganisation(user: User with Entity, organisation: Organisation with Entity, profile: Profile with Entity)(implicit
       graph: Graph,
@@ -64,21 +96,6 @@ class UserSrv @Inject() (
       } yield richUser
     }
 
-  def addOrCreateUser(user: User, avatar: Option[FFile], organisation: Organisation with Entity, profile: Profile with Entity)(implicit
-      graph: Graph,
-      authContext: AuthContext
-  ): Try[RichUser] =
-    getByName(user.login)
-      .getOrFail("User")
-      .orElse {
-        for {
-          validUser   <- checkUser(user)
-          createdUser <- createEntity(validUser)
-          _           <- avatar.map(setAvatar(createdUser, _)).flip
-        } yield createdUser
-      }
-      .flatMap(addUserToOrganisation(_, organisation, profile))
-
   def canSetPassword(user: User with Entity)(implicit graph: Graph, authContext: AuthContext): Boolean = {
     val userOrganisations     = get(user).organisations.value(_.name).toSet
     val operatorOrganisations = current.organisations(Permissions.manageUser).value(_.name).toSeq