#1365 Add case information in observable output

TheHive-Project · May 27, 2020 · 90aebf0 · 90aebf0
1 parent 4564863
commit 90aebf0
Show file tree

Hide file tree

Showing 5 changed files with 88 additions and 49 deletions.
diff --git a/dto/src/main/scala/org/thp/thehive/dto/v0/Observable.scala b/dto/src/main/scala/org/thp/thehive/dto/v0/Observable.scala
@@ -50,7 +50,8 @@ case class OutputObservable(
     message: Option[String],
     reports: JsObject,
     stats: JsObject,
-    seen: Option[Boolean]
+    seen: Option[Boolean],
+    `case`: Option[OutputCase]
 )
 
 object OutputObservable {

diff --git a/thehive/app/org/thp/thehive/controllers/v0/Conversion.scala b/thehive/app/org/thp/thehive/controllers/v0/Conversion.scala
@@ -364,40 +364,75 @@ object Conversion {
           )
       )
       .withFieldConst(_.stats, JsObject.empty)
+      .withFieldConst(_.`case`, None)
       .transform
   )
 
+  implicit val observableWithExtraOutput: Renderer.Aux[(RichObservable, JsObject, Option[RichCase]), OutputObservable] =
+    Renderer.json[(RichObservable, JsObject, Option[RichCase]), OutputObservable] {
+      case (richObservable, stats, richCase) =>
+        richObservable
+          .into[OutputObservable]
+          .withFieldConst(_._type, "case_artifact")
+          .withFieldComputed(_.id, _.observable._id)
+          .withFieldComputed(_._id, _.observable._id)
+          .withFieldComputed(_.updatedAt, _.observable._updatedAt)
+          .withFieldComputed(_.updatedBy, _.observable._updatedBy)
+          .withFieldComputed(_.createdAt, _.observable._createdAt)
+          .withFieldComputed(_.createdBy, _.observable._createdBy)
+          .withFieldComputed(_.dataType, _.`type`.name)
+          .withFieldComputed(_.startDate, _.observable._createdAt)
+          .withFieldComputed(_.tags, _.tags.map(_.toString).toSet)
+          .withFieldComputed(_.data, _.data.map(_.data))
+          .withFieldComputed(_.attachment, _.attachment.map(_.toValue))
+          .withFieldComputed(
+            _.reports, { a =>
+              JsObject(a.reportTags.groupBy(_.origin).map {
+                case (origin, tags) =>
+                  origin -> Json.obj(
+                    "taxonomies" -> tags
+                      .map(t => Json.obj("level" -> t.level.toString, "namespace" -> t.namespace, "predicate" -> t.predicate, "value" -> t.value))
+                  )
+              })
+            }
+          )
+          .withFieldConst(_.stats, stats)
+          .withFieldConst(_.`case`, richCase.map(_.toValue))
+          .transform
+    }
+
   implicit val observableWithStatsOutput: Renderer.Aux[(RichObservable, JsObject), OutputObservable] =
-    Renderer.json[(RichObservable, JsObject), OutputObservable](richObservableWithStats =>
-      richObservableWithStats
-        ._1
-        .into[OutputObservable]
-        .withFieldConst(_._type, "case_artifact")
-        .withFieldComputed(_.id, _.observable._id)
-        .withFieldComputed(_._id, _.observable._id)
-        .withFieldComputed(_.updatedAt, _.observable._updatedAt)
-        .withFieldComputed(_.updatedBy, _.observable._updatedBy)
-        .withFieldComputed(_.createdAt, _.observable._createdAt)
-        .withFieldComputed(_.createdBy, _.observable._createdBy)
-        .withFieldComputed(_.dataType, _.`type`.name)
-        .withFieldComputed(_.startDate, _.observable._createdAt)
-        .withFieldComputed(_.tags, _.tags.map(_.toString).toSet)
-        .withFieldComputed(_.data, _.data.map(_.data))
-        .withFieldComputed(_.attachment, _.attachment.map(_.toValue))
-        .withFieldComputed(
-          _.reports, { a =>
-            JsObject(a.reportTags.groupBy(_.origin).map {
-              case (origin, tags) =>
-                origin -> Json.obj(
-                  "taxonomies" -> tags
-                    .map(t => Json.obj("level" -> t.level.toString, "namespace" -> t.namespace, "predicate" -> t.predicate, "value" -> t.value))
-                )
-            })
-          }
-        )
-        .withFieldConst(_.stats, richObservableWithStats._2)
-        .transform
-    )
+    Renderer.json[(RichObservable, JsObject), OutputObservable] {
+      case (richObservable, stats) =>
+        richObservable
+          .into[OutputObservable]
+          .withFieldConst(_._type, "case_artifact")
+          .withFieldComputed(_.id, _.observable._id)
+          .withFieldComputed(_._id, _.observable._id)
+          .withFieldComputed(_.updatedAt, _.observable._updatedAt)
+          .withFieldComputed(_.updatedBy, _.observable._updatedBy)
+          .withFieldComputed(_.createdAt, _.observable._createdAt)
+          .withFieldComputed(_.createdBy, _.observable._createdBy)
+          .withFieldComputed(_.dataType, _.`type`.name)
+          .withFieldComputed(_.startDate, _.observable._createdAt)
+          .withFieldComputed(_.tags, _.tags.map(_.toString).toSet)
+          .withFieldComputed(_.data, _.data.map(_.data))
+          .withFieldComputed(_.attachment, _.attachment.map(_.toValue))
+          .withFieldComputed(
+            _.reports, { a =>
+              JsObject(a.reportTags.groupBy(_.origin).map {
+                case (origin, tags) =>
+                  origin -> Json.obj(
+                    "taxonomies" -> tags
+                      .map(t => Json.obj("level" -> t.level.toString, "namespace" -> t.namespace, "predicate" -> t.predicate, "value" -> t.value))
+                  )
+              })
+            }
+          )
+          .withFieldConst(_.stats, stats)
+          .withFieldConst(_.`case`, None)
+          .transform
+    }
 
   implicit class InputOrganisationOps(inputOrganisation: InputOrganisation) {
 

diff --git a/thehive/app/org/thp/thehive/controllers/v0/ObservableCtrl.scala b/thehive/app/org/thp/thehive/controllers/v0/ObservableCtrl.scala
@@ -39,22 +39,27 @@ class ObservableCtrl @Inject() (
     FieldsParser[IdOrName],
     (param, graph, authContext) => observableSrv.get(param.idOrName)(graph).visible(authContext)
   )
-  override val pageQuery: ParamQuery[OutputParam] = Query.withParam[OutputParam, ObservableSteps, PagedResult[(RichObservable, JsObject)]](
-    "page",
-    FieldsParser[OutputParam], {
-      case (OutputParam(from, to, withStats, _), observableSteps, authContext) =>
-        observableSteps
-          .richPage(from, to, withTotal = true) {
-            case o if withStats =>
-              o.richObservableWithCustomRenderer(observableStatsRenderer(authContext))
-            case o =>
-              o.richObservable.map(_ -> JsObject.empty)
-          }
-    }
-  )
+  override val pageQuery: ParamQuery[OutputParam] =
+    Query.withParam[OutputParam, ObservableSteps, PagedResult[(RichObservable, JsObject, Option[RichCase])]](
+      "page",
+      FieldsParser[OutputParam], {
+        case (OutputParam(from, to, withStats, 0), observableSteps, authContext) =>
+          observableSteps
+            .richPage(from, to, withTotal = true) {
+              case o if withStats =>
+                o.richObservableWithCustomRenderer(observableStatsRenderer(authContext)).map(ros => (ros._1, ros._2, None))
+              case o =>
+                o.richObservable.map(ro => (ro, JsObject.empty, None))
+            }
+        case (OutputParam(from, to, _, _), observableSteps, authContext) =>
+          observableSteps.richPage(from, to, withTotal = true)(
+            _.richObservableWithCustomRenderer(o => o.`case`.richCase(authContext)).map(roc => (roc._1, JsObject.empty, Some(roc._2)))
+          )
+      }
+    )
   override val outputQuery: Query = Query.output[RichObservable, ObservableSteps](_.richObservable)
   override val extraQueries: Seq[ParamQuery[_]] = Seq(
-    Query.output[(RichObservable, JsObject)]
+//    Query.output[(RichObservable, JsObject, Option[RichCase])]
   )
 
   def create(caseId: String): Action[AnyContent] =

diff --git a/thehive/app/org/thp/thehive/controllers/v0/TaskCtrl.scala b/thehive/app/org/thp/thehive/controllers/v0/TaskCtrl.scala
@@ -36,7 +36,7 @@ class TaskCtrl @Inject() (
     FieldsParser[OutputParam], {
       case (OutputParam(from, to, _, 0), taskSteps, _) => taskSteps.richPage(from, to, withTotal = true)(_.richTask.map(_ -> None))
       case (OutputParam(from, to, _, _), taskSteps, authContext) =>
-        taskSteps.richPage(from, to, withTotal = true)(_.richTaskWithCustomRenderer(_.`case`.richCase(authContext).map(c => Option(c)))(authContext))
+        taskSteps.richPage(from, to, withTotal = true)(_.richTaskWithCustomRenderer(_.`case`.richCase(authContext).map(c => Some(c))))
     }
   )
   override val getQuery: ParamQuery[IdOrName] = Query.initWithParam[IdOrName, TaskSteps](

diff --git a/thehive/app/org/thp/thehive/services/TaskSrv.scala b/thehive/app/org/thp/thehive/services/TaskSrv.scala
@@ -161,9 +161,7 @@ class TaskSteps(raw: GremlinScala[Vertex])(implicit db: Database, graph: Graph)
         }
     )
 
-  def richTaskWithCustomRenderer[A](
-      entityRenderer: TaskSteps => TraversalLike[_, A]
-  )(implicit authContext: AuthContext): Traversal[(RichTask, A), (RichTask, A)] =
+  def richTaskWithCustomRenderer[A](entityRenderer: TaskSteps => TraversalLike[_, A]): Traversal[(RichTask, A), (RichTask, A)] =
     Traversal(
       raw
         .project(