#1237 Prevent multiple links between user and org, even during migration

TheHive-Project · Mar 4, 2020 · b2e8a90 · b2e8a90
1 parent 8d26c37
commit b2e8a90
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 5 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
@@ -210,7 +210,8 @@ class Output @Inject() (
             for {
               organisation <- getOrganisation(organisationName)
               profile      <- profileSrv.getOrFail(profileName)
-            } yield roleSrv.create(createdUser, organisation, profile)
+              _            <- userSrv.add(createdUser, organisation, profile)
+            } yield ()
         }
       } yield IdMapping(inputUser.metaData.id, createdUser._id)
   }

diff --git a/thehive/app/org/thp/thehive/services/UserSrv.scala b/thehive/app/org/thp/thehive/services/UserSrv.scala
@@ -59,9 +59,11 @@ class UserSrv @Inject() (configuration: Configuration, roleSrv: RoleSrv, auditSr
       implicit graph: Graph,
       authContext: AuthContext
   ): Try[RichUser] =
-    get(user).richUser(organisation.name).getOrFail().orElse {
+    (if (!get(user).organisations.getByName(organisation.name).exists())
+       roleSrv.create(user, organisation, profile)
+     else
+       Success(())).flatMap { _ =>
       for {
-        _        <- roleSrv.create(user, organisation, profile)
         richUser <- get(user).richUser(organisation.name).getOrFail()
         _        <- auditSrv.user.create(user, richUser.toJson)
       } yield richUser
@@ -182,8 +184,8 @@ class UserSteps(raw: GremlinScala[Vertex])(implicit db: Database, graph: Graph)
     )
 
   def organisations(requiredPermission: String): OrganisationSteps = {
-    val isInDefaultOrganisation = newInstance().organisations0(requiredPermission).get(OrganisationSrv.administration.name).exists()
-    if (isInDefaultOrganisation) new OrganisationSteps(db.labelFilter(Model.vertex[Organisation])(graph.V))
+    val isInAdminOrganisation = newInstance().organisations0(requiredPermission).get(OrganisationSrv.administration.name).exists()
+    if (isInAdminOrganisation) new OrganisationSteps(db.labelFilter(Model.vertex[Organisation])(graph.V))
     else organisations0(requiredPermission)
   }