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 index 13f26c4906..089221d5ca 100644 --- 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 index 50bcbc8f7c..75ca024c34 100644 --- 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 index ea2275c5aa..0fa7e13bed 100644 --- 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 index 094fbfa373..73fd22ad2c 100644 --- 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 index ec8011232f..bdbe53ca14 100644 --- 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