TheHive-Project#2249 Fix message template

TKCERT · Nov 18, 2021 · 1eaaa6a · 1eaaa6a
1 parent 8aee43d
commit 1eaaa6a
Show file tree

Hide file tree

Showing 11 changed files with 55 additions and 34 deletions.
diff --git a/.../scala/org/thp/thehive/connector/cortex/services/notification/notifiers/RunAnalyzer.scala b/.../scala/org/thp/thehive/connector/cortex/services/notification/notifiers/RunAnalyzer.scala
@@ -64,8 +64,8 @@ class RunAnalyzer(
 
   override def execute(
       audit: Audit with Entity,
-      context: Option[Entity],
-      `object`: Option[Entity],
+      context: Option[Map[String, Seq[Any]] with Entity],
+      `object`: Option[Map[String, Seq[Any]] with Entity],
       organisation: Organisation with Entity,
       user: Option[User with Entity]
   )(implicit graph: Graph): Future[Unit] =

diff --git a/...scala/org/thp/thehive/connector/cortex/services/notification/notifiers/RunResponder.scala b/...scala/org/thp/thehive/connector/cortex/services/notification/notifiers/RunResponder.scala
@@ -82,8 +82,8 @@ class RunResponder(
 
   override def execute(
       audit: Audit with Entity,
-      context: Option[Entity],
-      `object`: Option[Entity],
+      context: Option[Map[String, Seq[Any]] with Entity],
+      `object`: Option[Map[String, Seq[Any]] with Entity],
       organisation: Organisation with Entity,
       user: Option[User with Entity]
   )(implicit graph: Graph): Future[Unit] =

diff --git a/thehive/app/org/thp/thehive/services/AuditSrv.scala b/thehive/app/org/thp/thehive/services/AuditSrv.scala
@@ -321,20 +321,23 @@ object AuditOps {
   }
 
   implicit class AuditOpsDefs(traversal: Traversal.V[Audit]) {
-    def auditContextObjectOrganisation
-        : Traversal[(Audit with Entity, Option[Entity], Option[Entity], Option[Organisation with Entity]), JMap[String, Any], Converter[
-          (Audit with Entity, Option[Entity], Option[Entity], Option[Organisation with Entity]),
-          JMap[String, Any]
-        ]] =
+    def auditContextObjectOrganisation: Traversal[
+      (Audit with Entity, Option[Map[String, Seq[Any]] with Entity], Option[Map[String, Seq[Any]] with Entity], Seq[Organisation with Entity]),
+      JMap[String, Any],
+      Converter[
+        (Audit with Entity, Option[Map[String, Seq[Any]] with Entity], Option[Map[String, Seq[Any]] with Entity], Seq[Organisation with Entity]),
+        JMap[String, Any]
+      ]
+    ] =
       traversal
         .project(
           _.by
-            .by(_.context.entity.fold)
-            .by(_.`object`.entity.fold)
+            .by(_.context.entityMap.option)
+            .by(_.`object`.entityMap.option)
             .by(_.organisation.v[Organisation].fold)
         )
         .domainMap {
-          case (audit, context, obj, organisation) => (audit, context.headOption, obj.headOption, organisation.headOption)
+          case (audit, context, obj, organisation) => (audit, context, obj, organisation)
         }
 
     def richAudit: Traversal[RichAudit, JMap[String, Any], Converter[RichAudit, JMap[String, Any]]] =

diff --git a/thehive/app/org/thp/thehive/services/notification/NotificationActor.scala b/thehive/app/org/thp/thehive/services/notification/NotificationActor.scala
@@ -114,8 +114,8 @@ class NotificationActor @Inject() (
       user: Option[User with Entity],
       notificationConfigs: Seq[NotificationConfig],
       audit: Audit with Entity,
-      context: Option[Entity],
-      `object`: Option[Entity],
+      context: Option[Map[String, Seq[Any]] with Entity],
+      `object`: Option[Map[String, Seq[Any]] with Entity],
       organisation: Organisation with Entity
   )(implicit
       graph: Graph

diff --git a/thehive/app/org/thp/thehive/services/notification/notifiers/AppendToFile.scala b/thehive/app/org/thp/thehive/services/notification/notifiers/AppendToFile.scala
@@ -37,8 +37,8 @@ class AppendToFile(filename: String, template: String, charset: Charset, baseUrl
 
   override def execute(
       audit: Audit with Entity,
-      context: Option[Entity],
-      `object`: Option[Entity],
+      context: Option[Map[String, Seq[Any]] with Entity],
+      `object`: Option[Map[String, Seq[Any]] with Entity],
       organisation: Organisation with Entity,
       user: Option[User with Entity]
   )(implicit

diff --git a/thehive/app/org/thp/thehive/services/notification/notifiers/Emailer.scala b/thehive/app/org/thp/thehive/services/notification/notifiers/Emailer.scala
@@ -50,8 +50,8 @@ class Emailer(
 
   override def execute(
       audit: Audit with Entity,
-      context: Option[Entity],
-      `object`: Option[Entity],
+      context: Option[Map[String, Seq[Any]] with Entity],
+      `object`: Option[Map[String, Seq[Any]] with Entity],
       organisation: Organisation with Entity,
       user: Option[User with Entity]
   )(implicit

diff --git a/thehive/app/org/thp/thehive/services/notification/notifiers/Mattermost.scala b/thehive/app/org/thp/thehive/services/notification/notifiers/Mattermost.scala
@@ -65,8 +65,8 @@ class Mattermost(ws: WSClient, mattermostNotification: MattermostNotification, b
 
   def execute(
       audit: Audit with Entity,
-      context: Option[Entity],
-      `object`: Option[Entity],
+      context: Option[Map[String, Seq[Any]] with Entity],
+      `object`: Option[Map[String, Seq[Any]] with Entity],
       organisation: Organisation with Entity,
       user: Option[User with Entity]
   )(implicit

diff --git a/thehive/app/org/thp/thehive/services/notification/notifiers/Notifier.scala b/thehive/app/org/thp/thehive/services/notification/notifiers/Notifier.scala
@@ -15,8 +15,8 @@ trait Notifier {
 
   def execute(
       audit: Audit with Entity,
-      context: Option[Entity],
-      `object`: Option[Entity],
+      context: Option[Map[String, Seq[Any]] with Entity],
+      `object`: Option[Map[String, Seq[Any]] with Entity],
       organisation: Organisation with Entity,
       user: Option[User with Entity]
   )(implicit

diff --git a/thehive/app/org/thp/thehive/services/notification/notifiers/Template.scala b/thehive/app/org/thp/thehive/services/notification/notifiers/Template.scala
@@ -2,7 +2,7 @@ package org.thp.thehive.services.notification.notifiers
 
 import com.github.jknack.handlebars.Handlebars
 import com.github.jknack.handlebars.helper.ConditionalHelpers
-import org.thp.scalligraph.models.{Entity, Schema}
+import org.thp.scalligraph.models.{Entity, MappingCardinality, Schema}
 import org.thp.thehive.models.{Audit, User}
 
 import java.util.{HashMap => JHashMap}
@@ -18,6 +18,24 @@ trait Template {
     * @param cc the entity
     * @return
     */
+  private def getMap(cc: Map[String, Seq[Any]] with Entity): Map[String, String] =
+    schema.getModel(cc._label).fold(cc.mapValues(_.mkString("[", ",", "]"))) { model =>
+      cc.map {
+        case (k, v) =>
+          model.fields.get(k).fold(k -> v.mkString("[", ",", "]")) {
+            case mapping if mapping.cardinality == MappingCardinality.list || mapping.cardinality == MappingCardinality.set =>
+              k -> v.mkString("[", ",", "]")
+            case _ => k -> v.head.toString
+          }
+      } +
+        ("_id"        -> cc._id.toString) +
+        ("_type"      -> cc._label) +
+        ("_createdAt" -> cc._createdAt.toString) +
+        ("_createdBy" -> cc._createdBy) +
+        ("_updatedAt" -> cc._updatedAt.fold("never")(_.toString)) +
+        ("_updatedBy" -> cc._updatedBy.getOrElse("nobody"))
+    }
+
   private def getMap(cc: Entity): Map[String, String] =
     schema
       .getModel(cc._label)
@@ -63,8 +81,8 @@ trait Template {
   def buildMessage(
       template: String,
       audit: Audit with Entity,
-      context: Option[Entity],
-      `object`: Option[Entity],
+      context: Option[Map[String, Seq[Any]] with Entity],
+      `object`: Option[Map[String, Seq[Any]] with Entity],
       user: Option[User with Entity],
       baseUrl: String
   ): Try[String] = {

diff --git a/thehive/app/org/thp/thehive/services/notification/notifiers/Webhook.scala b/thehive/app/org/thp/thehive/services/notification/notifiers/Webhook.scala
@@ -246,8 +246,8 @@ class Webhook(
 
   override def execute(
       audit: Audit with Entity,
-      context: Option[Entity],
-      `object`: Option[Entity],
+      context: Option[Map[String, Seq[Any]] with Entity],
+      `object`: Option[Map[String, Seq[Any]] with Entity],
       organisation: Organisation with Entity,
       user: Option[User with Entity]
   )(implicit graph: Graph): Future[Unit] =

diff --git a/thehive/test/org/thp/thehive/services/notification/notifiers/NotificationTemplateTest.scala b/thehive/test/org/thp/thehive/services/notification/notifiers/NotificationTemplateTest.scala
@@ -83,14 +83,14 @@ class NotificationTemplateTest extends PlaySpecification with TestAppBuilder {
           |Context {{context._id}}""".stripMargin
 
       val message = app[Database].tryTransaction { implicit graph =>
-        println(s"querying ${graph} ${graph.db} ${graph.db}")
         for {
-          case4 <- app[CaseSrv].get(EntityName("1")).getOrFail("Case")
-          _     <- app[CaseSrv].addTags(case4, Set("emailer test"))
-          _     <- app[CaseSrv].addTags(case4, Set("emailer test")) // this is needed to make AuditSrv write Audit in DB
-          audit <- app[AuditSrv].startTraversal.has(_.objectId, case4._id.toString).getOrFail("Audit")
-          user  <- app[UserSrv].get(EntityName("certuser@thehive.local")).getOrFail("User")
-          msg   <- templateEngine(app[Schema]).buildMessage(template, audit, Some(case4), Some(case4), Some(user), "http://localhost/")
+          case4       <- app[CaseSrv].get(EntityName("1")).getOrFail("Case")
+          case4Entity <- app[CaseSrv].get(EntityName("1")).entityMap.getOrFail("Case")
+          _           <- app[CaseSrv].addTags(case4, Set("emailer test"))
+          _           <- app[CaseSrv].addTags(case4, Set("emailer test")) // this is needed to make AuditSrv write Audit in DB
+          audit       <- app[AuditSrv].startTraversal.has(_.objectId, case4._id.toString).getOrFail("Audit")
+          user        <- app[UserSrv].get(EntityName("certuser@thehive.local")).getOrFail("User")
+          msg         <- templateEngine(app[Schema]).buildMessage(template, audit, Some(case4Entity), Some(case4Entity), Some(user), "http://localhost/")
         } yield msg
       }
       message must beSuccessfulTry.which { m =>