#1766 Add procedure query in case and extraData in procedure

TheHive-Project · Feb 11, 2021 · c66164f · c66164f
1 parent 0395890
commit c66164f
Show file tree

Hide file tree

Showing 6 changed files with 63 additions and 5 deletions.
diff --git a/dto/src/main/scala/org/thp/thehive/dto/v1/Procedure.scala b/dto/src/main/scala/org/thp/thehive/dto/v1/Procedure.scala
@@ -1,6 +1,6 @@
 package org.thp.thehive.dto.v1
 
-import play.api.libs.json.{Format, Json, Reads, Writes}
+import play.api.libs.json.{Format, JsObject, Json, Reads, Writes}
 
 import java.util.Date
 
@@ -37,7 +37,8 @@ case class OutputProcedure(
     _updatedBy: Option[String],
     description: String,
     occurence: Date,
-    patternId: String
+    patternId: String,
+    extraData: JsObject
 )
 
 object OutputProcedure {

diff --git a/thehive/app/org/thp/thehive/controllers/v1/CaseCtrl.scala b/thehive/app/org/thp/thehive/controllers/v1/CaseCtrl.scala
@@ -60,7 +60,8 @@ class CaseCtrl @Inject() (
     Query[Traversal.V[Case], Traversal.V[User]]("assignableUsers", (caseSteps, authContext) => caseSteps.assignableUsers(authContext)),
     Query[Traversal.V[Case], Traversal.V[Organisation]]("organisations", (caseSteps, authContext) => caseSteps.organisations.visible(authContext)),
     Query[Traversal.V[Case], Traversal.V[Alert]]("alerts", (caseSteps, authContext) => caseSteps.alert.visible(authContext)),
-    Query[Traversal.V[Case], Traversal.V[Share]]("shares", (caseSteps, authContext) => caseSteps.shares.visible(authContext))
+    Query[Traversal.V[Case], Traversal.V[Share]]("shares", (caseSteps, authContext) => caseSteps.shares.visible(authContext)),
+    Query[Traversal.V[Case], Traversal.V[Procedure]]("procedures", (caseSteps, _) => caseSteps.procedure)
   )
 
   def create: Action[AnyContent] =

diff --git a/thehive/app/org/thp/thehive/controllers/v1/Conversion.scala b/thehive/app/org/thp/thehive/controllers/v1/Conversion.scala
@@ -546,7 +546,19 @@ object Conversion {
       _.into[OutputProcedure]
         .withFieldComputed(_._id, _._id.toString)
         .withFieldComputed(_.patternId, _.pattern.patternId)
+        .withFieldConst(_.extraData, JsObject.empty)
         .transform
     )
 
+  implicit val richProcedureWithStatsRenderer: Renderer.Aux[(RichProcedure, JsObject), OutputProcedure] =
+    Renderer.toJson[(RichProcedure, JsObject), OutputProcedure] { procedureWithExtraData =>
+      procedureWithExtraData
+        ._1
+        .into[OutputProcedure]
+        .withFieldComputed(_._id, _._id.toString)
+        .withFieldComputed(_.patternId, _.pattern.patternId)
+        .withFieldConst(_.extraData, procedureWithExtraData._2)
+        .transform
+    }
+
 }
diff --git a/thehive/app/org/thp/thehive/controllers/v1/ProcedureCtrl.scala b/thehive/app/org/thp/thehive/controllers/v1/ProcedureCtrl.scala
@@ -21,7 +21,8 @@ class ProcedureCtrl @Inject() (
     properties: Properties,
     procedureSrv: ProcedureSrv,
     @Named("with-thehive-schema") implicit val db: Database
-) extends QueryableCtrl {
+) extends QueryableCtrl
+    with ProcedureRenderer {
   override val entityName: String                 = "procedure"
   override val publicProperties: PublicProperties = properties.procedure
   override val initialQuery: Query = Query.init[Traversal.V[Procedure]](
@@ -33,7 +34,10 @@ class ProcedureCtrl @Inject() (
   override val pageQuery: ParamQuery[OutputParam] = Query.withParam[OutputParam, Traversal.V[Procedure], IteratorOutput](
     "page",
     FieldsParser[OutputParam],
-    (range, procedureSteps, _) => procedureSteps.richPage(range.from, range.to, range.extraData.contains("total"))(_.richProcedure)
+    (range, procedureSteps, _) =>
+      procedureSteps.richPage(range.from, range.to, range.extraData.contains("total"))(
+        _.richProcedureWithCustomRenderer(procedureStatsRenderer(range.extraData - "total"))
+      )
   )
   override val outputQuery: Query = Query.output[RichProcedure, Traversal.V[Procedure]](_.richProcedure)
   override val getQuery: ParamQuery[EntityIdOrName] = Query.initWithParam[EntityIdOrName, Traversal.V[Procedure]](

diff --git a/thehive/app/org/thp/thehive/controllers/v1/ProcedureRenderer.scala b/thehive/app/org/thp/thehive/controllers/v1/ProcedureRenderer.scala
@@ -0,0 +1,27 @@
+package org.thp.thehive.controllers.v1
+
+import org.thp.scalligraph.traversal.{Converter, Traversal}
+import org.thp.thehive.controllers.v1.Conversion._
+import org.thp.thehive.models.Procedure
+import org.thp.thehive.services.PatternOps._
+import org.thp.thehive.services.ProcedureOps._
+import play.api.libs.json.JsValue
+
+import java.util.{Map => JMap}
+
+trait ProcedureRenderer extends BaseRenderer[Procedure] {
+  def patternStats: Traversal.V[Procedure] => Traversal[JsValue, JMap[String, Any], Converter[JsValue, JMap[String, Any]]] =
+    _.pattern.richPattern.domainMap(_.toJson)
+
+  def procedureStatsRenderer(extraData: Set[String]): Traversal.V[Procedure] => JsTraversal = { implicit traversal =>
+    baseRenderer(
+      extraData,
+      traversal,
+      {
+        case (f, "pattern") => addData("pattern", f)(patternStats)
+        case (f, _)         => f
+      }
+    )
+  }
+
+}
diff --git a/thehive/app/org/thp/thehive/services/ProcedureSrv.scala b/thehive/app/org/thp/thehive/services/ProcedureSrv.scala
@@ -80,5 +80,18 @@ object ProcedureOps {
             .by(_.pattern)
         )
         .domainMap { case (procedure, pattern) => RichProcedure(procedure, pattern) }
+
+    def richProcedureWithCustomRenderer[D, G, C <: Converter[D, G]](
+        entityRenderer: Traversal.V[Procedure] => Traversal[D, G, C]
+    ): Traversal[(RichProcedure, D), JMap[String, Any], Converter[(RichProcedure, D), JMap[String, Any]]] =
+      traversal
+        .project(
+          _.by
+            .by(_.pattern)
+            .by(entityRenderer)
+        )
+        .domainMap {
+          case (procedure, pattern, renderedEntity) => (RichProcedure(procedure, pattern), renderedEntity)
+        }
   }
 }