From fc9262d946a91553dd8a8f5c0dfde16e1405e36c Mon Sep 17 00:00:00 2001
From: To-om <toom@thehive-project.org>
Date: Wed, 22 Jul 2020 11:33:38 +0200
Subject: [PATCH] #1410 Fix organisation in user list

---
 .../thp/thehive/controllers/v0/UserCtrl.scala |  4 ++--
 .../thp/thehive/controllers/v1/UserCtrl.scala |  4 ++--
 .../org/thp/thehive/services/UserSrv.scala    | 24 ++++++++++++++++++-
 3 files changed, 27 insertions(+), 5 deletions(-)

diff --git a/thehive/app/org/thp/thehive/controllers/v0/UserCtrl.scala b/thehive/app/org/thp/thehive/controllers/v0/UserCtrl.scala
index dc8eb06ab4..33527af033 100644
--- a/thehive/app/org/thp/thehive/controllers/v0/UserCtrl.scala
+++ b/thehive/app/org/thp/thehive/controllers/v0/UserCtrl.scala
@@ -47,10 +47,10 @@ class UserCtrl @Inject() (
   override val pageQuery: ParamQuery[OutputParam] = Query.withParam[OutputParam, UserSteps, PagedResult[RichUser]](
     "page",
     FieldsParser[OutputParam],
-    (range, userSteps, authContext) => userSteps.richUser(authContext.organisation).page(range.from, range.to, withTotal = true)
+    (range, userSteps, authContext) => userSteps.richUser(authContext).page(range.from, range.to, withTotal = true)
   )
   override val outputQuery: Query =
-    Query.outputWithContext[RichUser, UserSteps]((userSteps, authContext) => userSteps.richUser(authContext.organisation))
+    Query.outputWithContext[RichUser, UserSteps]((userSteps, authContext) => userSteps.richUser(authContext))
 
   override val extraQueries: Seq[ParamQuery[_]] = Seq()
 
diff --git a/thehive/app/org/thp/thehive/controllers/v1/UserCtrl.scala b/thehive/app/org/thp/thehive/controllers/v1/UserCtrl.scala
index 5a1c9ca06b..0a158d4821 100644
--- a/thehive/app/org/thp/thehive/controllers/v1/UserCtrl.scala
+++ b/thehive/app/org/thp/thehive/controllers/v1/UserCtrl.scala
@@ -51,10 +51,10 @@ class UserCtrl @Inject() (
   override val pageQuery: ParamQuery[OutputParam] = Query.withParam[OutputParam, UserSteps, PagedResult[RichUser]](
     "page",
     FieldsParser[OutputParam],
-    (range, userSteps, authContext) => userSteps.richUser(authContext.organisation).page(range.from, range.to, range.extraData.contains("total"))
+    (range, userSteps, authContext) => userSteps.richUser(authContext).page(range.from, range.to, range.extraData.contains("total"))
   )
   override val outputQuery: Query =
-    Query.outputWithContext[RichUser, UserSteps]((userSteps, authContext) => userSteps.richUser(authContext.organisation))
+    Query.outputWithContext[RichUser, UserSteps]((userSteps, authContext) => userSteps.richUser(authContext))
 
   override val extraQueries: Seq[ParamQuery[_]] = Seq(
     Query.init[UserSteps]("currentUser", (graph, authContext) => userSrv.current(graph, authContext)),
diff --git a/thehive/app/org/thp/thehive/services/UserSrv.scala b/thehive/app/org/thp/thehive/services/UserSrv.scala
index 10394c40c0..09e82778b2 100644
--- a/thehive/app/org/thp/thehive/services/UserSrv.scala
+++ b/thehive/app/org/thp/thehive/services/UserSrv.scala
@@ -3,6 +3,7 @@ package org.thp.thehive.services
 import java.util.regex.Pattern
 import java.util.{List => JList}
 
+import scala.collection.JavaConverters._
 import akka.actor.ActorRef
 import gremlin.scala._
 import javax.inject.{Inject, Named, Singleton}
@@ -14,7 +15,7 @@ import org.thp.scalligraph.query.PropertyUpdater
 import org.thp.scalligraph.services._
 import org.thp.scalligraph.steps.StepsOps._
 import org.thp.scalligraph.steps.{Traversal, TraversalLike, VertexSteps}
-import org.thp.scalligraph.{AuthorizationError, BadRequestError, EntitySteps, RichOptionTry}
+import org.thp.scalligraph.{AuthorizationError, BadRequestError, EntitySteps, InternalError, RichOptionTry}
 import org.thp.thehive.controllers.v1.Conversion._
 import org.thp.thehive.models._
 import play.api.Configuration
@@ -266,6 +267,27 @@ class UserSteps(raw: GremlinScala[Vertex])(implicit @Named("with-thehive-schema"
           RichUser(user.as[User], avatar, "", Set.empty, organisation)
       }
 
+  def richUser(implicit authContext: AuthContext): Traversal[RichUser, RichUser] =
+    this
+      .project(
+        _.by
+          .by(_.avatar.fold)
+          .by(_.role.project(_.by(_.profile).by(_.organisation.name)).fold)
+      )
+      .map {
+        case (user, attachment, profileOrganisations) =>
+          val po = profileOrganisations.asScala.map {
+            case (profile, organisationName) => profile.as[Profile] -> organisationName
+          }
+          po.find(_._2 == authContext.organisation)
+            .orElse(po.headOption)
+            .fold(throw InternalError(s"")) {
+              case (profile, organisationName) =>
+                val avatar = atMostOneOf[Vertex](attachment).map(_.as[Attachment].attachmentId)
+                RichUser(user.as[User], avatar, profile.name, profile.permissions, organisationName)
+            }
+      }
+
   def richUserWithCustomRenderer[A](organisation: String, entityRenderer: UserSteps => TraversalLike[_, A])(
       implicit authContext: AuthContext
   ): Traversal[(RichUser, A), (RichUser, A)] =