#1410 Add observable count in alert

TheHive-Project · Jun 27, 2020 · 558dae9 · 558dae9
1 parent bcb7413
commit 558dae9
Show file tree

Hide file tree

Showing 5 changed files with 36 additions and 14 deletions.
diff --git a/dto/src/main/scala/org/thp/thehive/dto/v1/Alert.scala b/dto/src/main/scala/org/thp/thehive/dto/v1/Alert.scala
@@ -47,7 +47,8 @@ case class OutputAlert(
     read: Boolean,
     follow: Boolean,
     customFields: Set[OutputCustomFieldValue] = Set.empty,
-    caseTemplate: Option[String] = None
+    caseTemplate: Option[String] = None,
+    observableCount: Long
 )
 
 object OutputAlert {

diff --git a/thehive/app/org/thp/thehive/models/Alert.scala b/thehive/app/org/thp/thehive/models/Alert.scala
@@ -62,7 +62,8 @@ case class RichAlert(
     tags: Seq[Tag with Entity],
     customFields: Seq[RichCustomField],
     caseId: Option[String],
-    caseTemplate: Option[String]
+    caseTemplate: Option[String],
+    observableCount: Long
 ) {
   def _id: String                  = alert._id
   def _createdAt: Date             = alert._createdAt
@@ -92,7 +93,8 @@ object RichAlert {
       tags: Seq[Tag with Entity],
       customFields: Seq[RichCustomField],
       caseId: Option[String],
-      caseTemplate: Option[String]
+      caseTemplate: Option[String],
+      observableCount: Long
   ): RichAlert =
     alert
       .asInstanceOf[Alert]
@@ -103,5 +105,6 @@ object RichAlert {
       .withFieldConst(_.customFields, customFields)
       .withFieldConst(_.caseId, caseId)
       .withFieldConst(_.caseTemplate, caseTemplate)
+      .withFieldConst(_.observableCount, observableCount)
       .transform
 }
diff --git a/thehive/app/org/thp/thehive/services/AlertSrv.scala b/thehive/app/org/thp/thehive/services/AlertSrv.scala
@@ -1,5 +1,6 @@
 package org.thp.thehive.services
 
+import java.lang.{Long => JLong}
 import java.util.{Date, List => JList}
 
 import gremlin.scala._
@@ -76,7 +77,7 @@ class AlertSrv @Inject() (
         _            <- caseTemplate.map(ct => alertCaseTemplateSrv.create(AlertCaseTemplate(), createdAlert, ct)).flip
         _            <- tags.toTry(t => alertTagSrv.create(AlertTag(), createdAlert, t))
         cfs          <- customFields.toTry { case (name, value) => createCustomField(createdAlert, name, value) }
-        richAlert = RichAlert(createdAlert, organisation.name, tags, cfs, None, caseTemplate.map(_.name))
+        richAlert = RichAlert(createdAlert, organisation.name, tags, cfs, None, caseTemplate.map(_.name), 0)
         _ <- auditSrv.alert.create(createdAlert, richAlert.toJson)
       } yield richAlert
   }
@@ -336,6 +337,7 @@ class AlertSteps(raw: GremlinScala[Vertex])(implicit @Named("with-thehive-schema
     val customFieldLabel      = StepLabel[JList[Path]]()
     val caseIdLabel           = StepLabel[JList[AnyRef]]()
     val caseTemplateNameLabel = StepLabel[JList[String]]()
+    val observableCountLabel  = StepLabel[JLong]()
     Traversal(
       raw
         .`match`(
@@ -348,9 +350,18 @@ class AlertSteps(raw: GremlinScala[Vertex])(implicit @Named("with-thehive-schema
             .has(Key("login") of authContext.userId),
           _.as(alertLabel).outToE[AlertCustomField].inV().path.fold.as(customFieldLabel),
           _.as(alertLabel).outTo[AlertCase].id().fold.as(caseIdLabel),
-          _.as(alertLabel).outTo[AlertCaseTemplate].values[String]("name").fold.as(caseTemplateNameLabel)
+          _.as(alertLabel).outTo[AlertCaseTemplate].values[String]("name").fold.as(caseTemplateNameLabel),
+          _.as(alertLabel).outToE[AlertObservable].count().as(observableCountLabel)
+        )
+        .select(
+          alertLabel.name,
+          organisationLabel.name,
+          tagLabel.name,
+          customFieldLabel.name,
+          caseIdLabel.name,
+          caseTemplateNameLabel.name,
+          observableCountLabel.name
         )
-        .select(alertLabel.name, organisationLabel.name, tagLabel.name, customFieldLabel.name, caseIdLabel.name, caseTemplateNameLabel.name)
         .map { resultMap =>
           val organisation = resultMap.getValue(organisationLabel).as[Organisation]
           val tags         = resultMap.getValue(tagLabel).asScala.map(_.as[Tag])
@@ -369,7 +380,8 @@ class AlertSteps(raw: GremlinScala[Vertex])(implicit @Named("with-thehive-schema
             tags,
             customFieldValues,
             atMostOneOf(resultMap.getValue(caseIdLabel)).map(_.toString),
-            atMostOneOf(resultMap.getValue(caseTemplateNameLabel))
+            atMostOneOf(resultMap.getValue(caseTemplateNameLabel)),
+            resultMap.getValue(observableCountLabel)
           ) -> organisation
         }
     )
@@ -395,10 +407,11 @@ class AlertSteps(raw: GremlinScala[Vertex])(implicit @Named("with-thehive-schema
             .and(By(__[Vertex].outToE[AlertCustomField].inV().path.fold))
             .and(By(__[Vertex].outTo[AlertCase].id().fold))
             .and(By(__[Vertex].outTo[AlertCaseTemplate].values[String]("name").fold))
+            .and(By(__[Vertex].outToE[AlertObservable].count()))
             .and(By(entityRenderer(newInstance(__[Vertex])).raw))
         )
         .map {
-          case (alert, organisation, tags, customFields, caseId, caseTemplate, renderedEntity) =>
+          case (alert, organisation, tags, customFields, caseId, caseTemplate, observableCount, renderedEntity) =>
             val customFieldValues = (customFields: JList[Path])
               .asScala
               .map(_.asScala.takeRight(2).toList.asInstanceOf[List[Element]])
@@ -412,7 +425,8 @@ class AlertSteps(raw: GremlinScala[Vertex])(implicit @Named("with-thehive-schema
               tags.asScala.map(_.as[Tag]),
               customFieldValues,
               atMostOneOf[AnyRef](caseId).map(_.toString),
-              atMostOneOf[String](caseTemplate)
+              atMostOneOf[String](caseTemplate),
+              observableCount
             ) -> renderedEntity
         }
     )
@@ -427,9 +441,10 @@ class AlertSteps(raw: GremlinScala[Vertex])(implicit @Named("with-thehive-schema
             .and(By(__[Vertex].outToE[AlertCustomField].inV().path.fold))
             .and(By(__[Vertex].outTo[AlertCase].id().fold))
             .and(By(__[Vertex].outTo[AlertCaseTemplate].values[String]("name").fold))
+            .and(By(__[Vertex].outToE[AlertObservable].count()))
         )
         .map {
-          case (alert, organisation, tags, customFields, caseId, caseTemplate) =>
+          case (alert, organisation, tags, customFields, caseId, caseTemplate, observableCount) =>
             val customFieldValues = (customFields: JList[Path])
               .asScala
               .map(_.asScala.takeRight(2).toList.asInstanceOf[List[Element]])
@@ -443,7 +458,8 @@ class AlertSteps(raw: GremlinScala[Vertex])(implicit @Named("with-thehive-schema
               tags.asScala.map(_.as[Tag]),
               customFieldValues,
               atMostOneOf[AnyRef](caseId).map(_.toString),
-              atMostOneOf[String](caseTemplate)
+              atMostOneOf[String](caseTemplate),
+              observableCount
             )
         }
     )

diff --git a/thehive/app/org/thp/thehive/services/ObservableSrv.scala b/thehive/app/org/thp/thehive/services/ObservableSrv.scala
@@ -36,7 +36,6 @@ class ObservableSrv @Inject() (
   val observableDataSrv        = new EdgeSrv[ObservableData, Observable, Data]
   val observableObservableType = new EdgeSrv[ObservableObservableType, Observable, ObservableType]
   val observableAttachmentSrv  = new EdgeSrv[ObservableAttachment, Observable, Attachment]
-  val alertObservableSrv       = new EdgeSrv[AlertObservable, Alert, Observable]
   val observableTagSrv         = new EdgeSrv[ObservableTag, Observable, Tag]
 
   override def steps(raw: GremlinScala[Vertex])(implicit graph: Graph): ObservableSteps = new ObservableSteps(raw)

diff --git a/thehive/test/org/thp/thehive/controllers/v1/AlertCtrlTest.scala b/thehive/test/org/thp/thehive/controllers/v1/AlertCtrlTest.scala
@@ -75,7 +75,9 @@ class AlertCtrlTest extends PlaySpecification with TestAppBuilder {
         pap = 2,
         read = false,
         follow = true,
-        customFields = Set.empty
+        customFields = Set.empty,
+        caseTemplate = None,
+        observableCount = 0L
       )
 
       createdAlert must_=== expected
@@ -120,7 +122,8 @@ class AlertCtrlTest extends PlaySpecification with TestAppBuilder {
         read = false,
         follow = true,
         customFields = Set.empty,
-        caseTemplate = Some("spam")
+        caseTemplate = Some("spam"),
+        observableCount = 0L
       )
 
       createdAlert must_=== expected