#1410 Add "isOwner" and "shareCount" extraData in observable and task

TheHive-Project · Jul 16, 2020 · a8ebc74 · a8ebc74
1 parent 73ee961
commit a8ebc74
Show file tree

Hide file tree

Showing 4 changed files with 27 additions and 5 deletions.
diff --git a/thehive/app/org/thp/thehive/controllers/v1/ObservableRenderer.scala b/thehive/app/org/thp/thehive/controllers/v1/ObservableRenderer.scala
@@ -15,7 +15,7 @@ import scala.collection.JavaConverters._
 
 trait ObservableRenderer {
 
-  def seenStats(observableSteps: ObservableSteps)(implicit authContext: AuthContext): Traversal[JsValue, JsValue] =
+  def seen(observableSteps: ObservableSteps)(implicit authContext: AuthContext): Traversal[JsValue, JsValue] =
     observableSteps
       .similar
       .visible
@@ -30,9 +30,17 @@ trait ObservableRenderer {
         )
       }
 
-  def sharesStats(observableSteps: ObservableSteps): Traversal[JsValue, JsValue] =
+  def shares(observableSteps: ObservableSteps): Traversal[JsValue, JsValue] =
     observableSteps.shares.organisation.name.fold.map(orgs => Json.toJson(orgs.asScala))
 
+  def shareCount(observableSteps: ObservableSteps): Traversal[JsValue, JsValue] =
+    observableSteps.organisations.count.map(count => JsNumber.apply(count.longValue()))
+
+  def isOwner(
+      observableSteps: ObservableSteps
+  )(implicit authContext: AuthContext): Traversal[JsValue, JsValue] =
+    observableSteps.origin.name.map(orgName => JsBoolean(orgName == authContext.organisation))
+
   def observableLinks(observableSteps: ObservableSteps): Traversal[JsValue, JsValue] =
     observableSteps.coalesce(
       _.alert.richAlert.map(a => Json.obj("alert"            -> a.toJson)),
@@ -55,8 +63,10 @@ trait ObservableRenderer {
       val dataName = extraData.toSeq
       dataName
         .foldLeft[ObservableSteps => GremlinScala[JMap[String, JsValue]]](_.raw.project(dataName.head, dataName.tail: _*)) {
-          case (f, "seen")        => f.andThen(addData(seenStats))
-          case (f, "shares")      => f.andThen(addData(sharesStats))
+          case (f, "seen")        => f.andThen(addData(seen))
+          case (f, "shares")      => f.andThen(addData(shares))
+          case (f, "isOwner")     => f.andThen(addData(isOwner))
+          case (f, "shareCount")  => f.andThen(addData(shareCount))
           case (f, "links")       => f.andThen(addData(observableLinks))
           case (f, "permissions") => f.andThen(addData(permissions))
           case (f, _)             => f.andThen(_.by(__.constant(JsNull).traversal))

diff --git a/thehive/app/org/thp/thehive/controllers/v1/TaskRenderer.scala b/thehive/app/org/thp/thehive/controllers/v1/TaskRenderer.scala
@@ -9,7 +9,7 @@ import org.thp.scalligraph.steps.StepsOps._
 import org.thp.scalligraph.steps.Traversal
 import org.thp.thehive.controllers.v1.Conversion._
 import org.thp.thehive.services.TaskSteps
-import play.api.libs.json.{JsNull, JsObject, JsString, JsValue}
+import play.api.libs.json._
 
 import scala.collection.JavaConverters._
 
@@ -27,6 +27,12 @@ trait TaskRenderer {
   def caseTemplateParentId(taskSteps: TaskSteps): Traversal[JsValue, JsValue] =
     taskSteps.caseTemplate.fold.map(_.asScala.headOption.fold[JsValue](JsNull)(ct => JsString(ct.id().toString)))
 
+  def shareCount(taskSteps: TaskSteps): Traversal[JsValue, JsValue] =
+    taskSteps.organisations.count.map(count => JsNumber.apply(count.longValue()))
+
+  def isOwner(taskSteps: TaskSteps)(implicit authContext: AuthContext): Traversal[JsValue, JsValue] =
+    taskSteps.origin.name.map(orgName => JsBoolean(orgName == authContext.organisation))
+
   def taskStatsRenderer(extraData: Set[String])(
       implicit authContext: AuthContext,
       db: Database,
@@ -44,6 +50,8 @@ trait TaskRenderer {
           case (f, "caseId")         => f.andThen(addData(caseParentId))
           case (f, "caseTemplate")   => f.andThen(addData(caseTemplateParent))
           case (f, "caseTemplateId") => f.andThen(addData(caseTemplateParentId))
+          case (f, "isOwner")        => f.andThen(addData(isOwner))
+          case (f, "shareCount")     => f.andThen(addData(shareCount))
           case (f, _)                => f.andThen(_.by(__.constant(JsNull).traversal))
         }
         .andThen(f => Traversal(f.map(m => JsObject(m.asScala))))

diff --git a/thehive/app/org/thp/thehive/services/ObservableSrv.scala b/thehive/app/org/thp/thehive/services/ObservableSrv.scala
@@ -232,6 +232,8 @@ class ObservableSteps(raw: GremlinScala[Vertex])(implicit @Named("with-thehive-s
 
   def organisations = new OrganisationSteps(raw.inTo[ShareObservable].inTo[OrganisationShare])
 
+  def origin: OrganisationSteps = new OrganisationSteps(raw.inTo[ShareCase].has(Key("owner") of true).inTo[OrganisationShare])
+
   override def newInstance(): ObservableSteps = new ObservableSteps(raw.clone())
 
   def richObservable: Traversal[RichObservable, RichObservable] =

diff --git a/thehive/app/org/thp/thehive/services/TaskSrv.scala b/thehive/app/org/thp/thehive/services/TaskSrv.scala
@@ -159,6 +159,8 @@ class TaskSteps(raw: GremlinScala[Vertex])(implicit db: Database, graph: Graph)
   def organisations(permission: Permission) =
     new OrganisationSteps(raw.inTo[ShareTask].filter(_.outTo[ShareProfile].has(Key("permissions") of permission)).inTo[OrganisationShare])
 
+  def origin: OrganisationSteps = new OrganisationSteps(raw.inTo[ShareCase].has(Key("owner") of true).inTo[OrganisationShare])
+
   def assignableUsers(implicit authContext: AuthContext): UserSteps =
     organisations(Permissions.manageTask)
       .visible