#1766 Add case and pattern data in stream when operation is related t…

…o a procedure
TheHive-Project · Mar 9, 2021 · cbe5f3d · cbe5f3d
1 parent 2a52b3b
commit cbe5f3d
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 15 deletions.
diff --git a/thehive/app/org/thp/thehive/controllers/v0/AuditRenderer.scala b/thehive/app/org/thp/thehive/controllers/v0/AuditRenderer.scala
@@ -10,41 +10,43 @@ import org.thp.thehive.services.AlertOps._
 import org.thp.thehive.services.AuditOps._
 import org.thp.thehive.services.CaseOps._
 import org.thp.thehive.services.LogOps._
+import org.thp.thehive.services.ProcedureOps._
 import org.thp.thehive.services.ObservableOps._
+import org.thp.thehive.services.PatternOps._
 import org.thp.thehive.services.TaskOps._
 import org.thp.thehive.services._
 import play.api.libs.json.{JsNumber, JsObject, JsString}
-
-import java.util.{Date, Map => JMap}
+import org.thp.thehive.controllers.v1.Conversion.{patternOutput, richProcedureRenderer}
+import java.util.{Date, List => JList, Map => JMap}
 
 trait AuditRenderer {
 
-  def caseToJson: Traversal.V[Case] => Traversal[JsObject, JMap[String, Any], Converter[JsObject, JMap[String, Any]]] =
-    _.richCaseWithoutPerms.domainMap[JsObject](_.toJson.as[JsObject])
+  def caseToJson: Traversal.V[Case] => Traversal[JsObject, JList[JMap[String, Any]], Converter[JsObject, JList[JMap[String, Any]]]] =
+    _.richCaseWithoutPerms.option.domainMap[JsObject](_.fold(JsObject.empty)(_.toJson.as[JsObject]))
 
   def taskToJson: Traversal.V[Task] => Traversal[JsObject, JMap[String, Any], Converter[JsObject, JMap[String, Any]]] =
     _.project(
-      _.by(_.richTaskWithoutActionRequired.domainMap(_.toJson))
+      _.by(_.richTaskWithoutActionRequired.option.domainMap(_.fold(JsObject.empty)(_.toJson.as[JsObject])))
         .by(t => caseToJson(t.`case`))
     ).domainMap {
-      case (task, case0) => task.as[JsObject] + ("case" -> case0)
+      case (task, case0) => task + ("case" -> case0)
     }
 
-  def alertToJson: Traversal.V[Alert] => Traversal[JsObject, JMap[String, Any], Converter[JsObject, JMap[String, Any]]] =
-    _.richAlert.domainMap(_.toJson.as[JsObject])
+  def alertToJson: Traversal.V[Alert] => Traversal[JsObject, JList[JMap[String, Any]], Converter[JsObject, JList[JMap[String, Any]]]] =
+    _.richAlert.option.domainMap(_.fold(JsObject.empty)(_.toJson.as[JsObject]))
 
   def logToJson: Traversal.V[Log] => Traversal[JsObject, JMap[String, Any], Converter[JsObject, JMap[String, Any]]] =
     _.project(
-      _.by(_.richLog.domainMap(_.toJson))
+      _.by(_.richLog.option.domainMap(_.fold(JsObject.empty)(_.toJson.as[JsObject])))
         .by(l => taskToJson(l.task))
-    ).domainMap { case (log, task) => log.as[JsObject] + ("case_task" -> task) }
+    ).domainMap { case (log, task) => log + ("case_task" -> task) }
 
   def observableToJson: Traversal.V[Observable] => Traversal[JsObject, JMap[String, Any], Converter[JsObject, JMap[String, Any]]] =
     _.project(
-      _.by(_.richObservable.domainMap(_.toJson))
+      _.by(_.richObservable.option.domainMap(_.fold(JsObject.empty)(_.toJson.as[JsObject])))
         .by(_.coalesceMulti(o => caseToJson(o.`case`), o => alertToJson(o.alert)))
     ).domainMap {
-      case (obs, caseOrAlert) => obs.as[JsObject] + ((caseOrAlert \ "_type").asOpt[String].getOrElse("<unknwon>") -> caseOrAlert)
+      case (obs, caseOrAlert) => obs + ((caseOrAlert \ "_type").asOpt[String].getOrElse("<unknown>") -> caseOrAlert)
     }
 
   case class Job(
@@ -81,6 +83,19 @@ trait AuditRenderer {
           )
       }
 
+  def patternToJson: Traversal.V[Pattern] => Traversal[JsObject, JList[JMap[String, Any]], Converter[JsObject, JList[JMap[String, Any]]]] =
+    _.richPattern.option.domainMap(_.fold(JsObject.empty)(_.toJson.as[JsObject]))
+
+  def procedureToJson: Traversal.V[Procedure] => Traversal[JsObject, JMap[String, Any], Converter[JsObject, JMap[String, Any]]] =
+    _.project(
+      _.by(_.richProcedure.option.domainMap(_.fold(JsObject.empty)(_.toJson.as[JsObject])))
+        .by(t => caseToJson(t.`case`))
+        .by(t => patternToJson(t.pattern))
+    )
+      .domainMap {
+        case (procedure, case0, pattern) => procedure + ("case" -> case0) + ("pattern" -> pattern)
+      }
+
   def auditRenderer: Traversal.V[Audit] => Traversal[JsObject, JMap[String, Any], Converter[JsObject, JMap[String, Any]]] =
     (_: Traversal.V[Audit])
       .coalesceIdent[Vertex](_.`object`, _.identity)
@@ -91,6 +106,7 @@ trait AuditRenderer {
           .option("Log", t => logToJson(t.v[Log]))
           .option("Observable", t => observableToJson(t.v[Observable]))
           .option("Alert", t => alertToJson(t.v[Alert]))
+          .option("Procedure", t => procedureToJson(t.v[Procedure]))
           .option("Job", jobToJson)
           .none(_.constant2[JsObject, JMap[String, Any]](JsObject.empty))
       )

diff --git a/thehive/app/org/thp/thehive/services/ProcedureSrv.scala b/thehive/app/org/thp/thehive/services/ProcedureSrv.scala
@@ -45,14 +45,14 @@ class ProcedureSrv @Inject() (
   )(implicit graph: Graph, authContext: AuthContext): Try[(Traversal.V[Procedure], JsObject)] =
     auditSrv.mergeAudits(super.update(traversal, propertyUpdaters)) {
       case (procedureSteps, updatedFields) =>
-        procedureSteps.clone().project(_.by.by(_.caze)).getOrFail("Procedure").flatMap {
+        procedureSteps.clone().project(_.by.by(_.`case`)).getOrFail("Procedure").flatMap {
           case (procedure, caze) => auditSrv.procedure.update(procedure, caze, updatedFields)
         }
     }
 
   def remove(procedure: Procedure with Entity)(implicit graph: Graph, authContext: AuthContext): Try[Unit] =
     for {
-      caze <- get(procedure).caze.getOrFail("Case")
+      caze <- get(procedure).`case`.getOrFail("Case")
       _    <- auditSrv.procedure.delete(procedure, caze)
     } yield get(procedure).remove()
 
@@ -64,7 +64,7 @@ object ProcedureOps {
     def pattern: Traversal.V[Pattern] =
       traversal.out[ProcedurePattern].v[Pattern]
 
-    def caze: Traversal.V[Case] =
+    def `case`: Traversal.V[Case] =
       traversal.in[CaseProcedure].v[Case]
 
     def get(idOrName: EntityIdOrName): Traversal.V[Procedure] =