#1731 Optimize observable queries for the use of the index

TheHive-Project · Jan 5, 2021 · 8a376e5 · 8a376e5
1 parent 3991fd0
commit 8a376e5
Show file tree

Hide file tree

Showing 19 changed files with 277 additions and 106 deletions.
diff --git a/...nnector/src/main/scala/org/thp/thehive/connector/cortex/services/ActionOperationSrv.scala b/...nnector/src/main/scala/org/thp/thehive/connector/cortex/services/ActionOperationSrv.scala
@@ -94,10 +94,19 @@ class ActionOperationSrv @Inject() (
 
       case AddArtifactToCase(_, dataType, dataMessage) =>
         for {
-          c       <- relatedCase.fold[Try[Case with Entity]](Failure(InternalError("Unable to apply action AddArtifactToCase without case")))(Success(_))
-          obsType <- observableTypeSrv.getOrFail(EntityIdOrName(dataType))
+          c            <- relatedCase.fold[Try[Case with Entity]](Failure(InternalError("Unable to apply action AddArtifactToCase without case")))(Success(_))
+          obsType      <- observableTypeSrv.getOrFail(EntityIdOrName(dataType))
+          organisation <- organisationSrv.getOrFail(authContext.organisation)
           richObservable <- observableSrv.create(
-            Observable(Some(dataMessage), 2, ioc = false, sighted = false, ignoreSimilarity = None),
+            Observable(
+              Some(dataMessage),
+              2,
+              ioc = false,
+              sighted = false,
+              ignoreSimilarity = None,
+              organisationIds = Seq(organisation._id),
+              relatedId = c._id
+            ),
             obsType,
             dataMessage,
             Set.empty[String],

diff --git a/cortex/connector/src/main/scala/org/thp/thehive/connector/cortex/services/Conversion.scala b/cortex/connector/src/main/scala/org/thp/thehive/connector/cortex/services/Conversion.scala
@@ -2,6 +2,7 @@ package org.thp.thehive.connector.cortex.services
 
 import io.scalaland.chimney.dsl._
 import org.thp.cortex.dto.v0.{OutputArtifact, OutputMinireport, JobStatus => CortexJobStatus}
+import org.thp.scalligraph.EntityId
 import org.thp.thehive.connector.cortex.models.JobStatus
 import org.thp.thehive.models.{Observable, ReportTag, ReportTagLevel}
 
@@ -21,14 +22,16 @@ object Conversion {
 
   implicit class CortexOutputArtifactOps(artifact: OutputArtifact) {
 
-    def toObservable: Observable =
+    def toObservable(relatedId: EntityId, organisationIds: EntityId*): Observable =
       artifact
         .into[Observable]
         .withFieldComputed(_.message, _.message)
         .withFieldComputed(_.tlp, _.tlp)
         .withFieldConst(_.ioc, false)
         .withFieldConst(_.sighted, false)
         .withFieldConst(_.ignoreSimilarity, None)
+        .withFieldConst(_.organisationIds, organisationIds)
+        .withFieldConst(_.relatedId, relatedId)
         .transform
   }
 

diff --git a/cortex/connector/src/main/scala/org/thp/thehive/connector/cortex/services/JobSrv.scala b/cortex/connector/src/main/scala/org/thp/thehive/connector/cortex/services/JobSrv.scala
@@ -212,11 +212,17 @@ class JobSrv @Inject() (
             Future
               .fromTry {
                 db.tryTransaction { implicit graph =>
-                  observableSrv
-                    .create(artifact.toObservable, dataType, artifact.data.get, artifact.tags, Nil)
-                    .flatMap { richObservable =>
-                      addObservable(job, richObservable.observable)
-                    }
+                  for {
+                    origObs <- get(job).observable.getOrFail("Observable")
+                    obs <- observableSrv.create(
+                      artifact.toObservable(job._id, origObs.organisationIds: _*),
+                      dataType,
+                      artifact.data.get,
+                      artifact.tags,
+                      Nil
+                    )
+                    _ <- addObservable(job, obs.observable)
+                  } yield ()
                 }
               }
           case Failure(e) => Future.failed(e)
@@ -258,9 +264,12 @@ class JobSrv @Inject() (
           savedAttachment <- Future.fromTry {
             db.tryTransaction { implicit graph =>
               for {
+                origObs           <- get(job).observable.getOrFail("Observable")
                 createdAttachment <- attachmentSrv.create(fFile)
-                richObservable    <- observableSrv.create(artifact.toObservable, attachmentType, createdAttachment, artifact.tags, Nil)
-                _                 <- reportObservableSrv.create(ReportObservable(), job, richObservable.observable)
+                richObservable <-
+                  observableSrv
+                    .create(artifact.toObservable(job._id, origObs.organisationIds: _*), attachmentType, createdAttachment, artifact.tags, Nil)
+                _ <- reportObservableSrv.create(ReportObservable(), job, richObservable.observable)
               } yield createdAttachment
             }
           }

diff --git a/migration/src/main/scala/org/thp/thehive/migration/th3/Conversion.scala b/migration/src/main/scala/org/thp/thehive/migration/th3/Conversion.scala
@@ -109,7 +109,7 @@ trait Conversion {
           )
     } yield InputObservable(
       metaData,
-      Observable(message, tlp, ioc, sighted, None),
+      Observable(message, tlp, ioc, sighted, None, Nil, EntityId("")), // organisation and related Ids are filled by output
       Seq(mainOrganisation),
       dataType,
       tags,
@@ -231,7 +231,15 @@ trait Conversion {
             )
       } yield InputObservable(
         metaData,
-        Observable(message, tlp.getOrElse(2), ioc.getOrElse(false), sighted = false, ignoreSimilarity = None),
+        Observable(
+          message,
+          tlp.getOrElse(2),
+          ioc.getOrElse(false),
+          sighted = false,
+          ignoreSimilarity = None,
+          organisationIds = Nil,
+          relatedId = EntityId("")
+        ),
         Nil,
         dataType,
         tags,
@@ -450,7 +458,7 @@ trait Conversion {
           )
       } yield InputObservable(
         metaData,
-        Observable(message, tlp, ioc, sighted, ignoreSimilarity = None),
+        Observable(message, tlp, ioc, sighted, ignoreSimilarity = None, organisationIds = Nil, relatedId = EntityId("")),
         Seq(mainOrganisation),
         dataType,
         tags,

diff --git a/migration/src/main/scala/org/thp/thehive/migration/th4/Output.scala b/migration/src/main/scala/org/thp/thehive/migration/th4/Output.scala
@@ -23,6 +23,7 @@ import org.thp.thehive.migration.IdMapping
 import org.thp.thehive.migration.dto._
 import org.thp.thehive.models._
 import org.thp.thehive.services._
+import org.thp.thehive.connector.cortex.services.JobOps._
 import play.api.cache.SyncCacheApi
 import play.api.cache.ehcache.EhCacheModule
 import play.api.inject.guice.GuiceInjector
@@ -600,23 +601,30 @@ class Output @Inject() (
       for {
         observableType <- getObservableType(inputObservable.`type`)
         tags           <- inputObservable.tags.filterNot(_.isEmpty).toTry(getTag)
+        orgs           <- inputObservable.organisations.toTry(getOrganisation)
         richObservable <-
           inputObservable
             .dataOrAttachment
             .fold(
               dataValue =>
                 dataSrv.createEntity(Data(dataValue)).flatMap { data =>
-                  observableSrv.create(inputObservable.observable, observableType, data, tags, Nil)
+                  observableSrv
+                    .create(inputObservable.observable.copy(organisationIds = orgs.map(_._id), relatedId = caseId), observableType, data, tags, Nil)
                 },
               inputAttachment =>
                 attachmentSrv.create(inputAttachment.name, inputAttachment.size, inputAttachment.contentType, inputAttachment.data).flatMap {
                   attachment =>
-                    observableSrv.create(inputObservable.observable, observableType, attachment, tags, Nil)
+                    observableSrv.create(
+                      inputObservable.observable.copy(organisationIds = orgs.map(_._id), relatedId = caseId),
+                      observableType,
+                      attachment,
+                      tags,
+                      Nil
+                    )
                 }
             )
         _ = updateMetaData(richObservable.observable, inputObservable.metaData)
         case0 <- getCase(caseId)
-        orgs  <- inputObservable.organisations.toTry(getOrganisation)
         _     <- orgs.toTry(o => shareSrv.shareObservable(richObservable, case0, o))
       } yield IdMapping(inputObservable.metaData.id, richObservable._id)
     }
@@ -636,6 +644,7 @@ class Output @Inject() (
       logger.debug(s"Create observable ${inputObservable.dataOrAttachment.fold(identity, _.name)} in job $jobId")
       for {
         job            <- jobSrv.getOrFail(jobId)
+        jobObs         <- jobSrv.get(job).observable.getOrFail("Observable")
         observableType <- getObservableType(inputObservable.`type`)
         tags = inputObservable.tags.filterNot(_.isEmpty).flatMap(getTag(_).toOption).toSeq
         richObservable <-
@@ -644,12 +653,25 @@ class Output @Inject() (
             .fold(
               dataValue =>
                 dataSrv.createEntity(Data(dataValue)).flatMap { data =>
-                  observableSrv.create(inputObservable.observable, observableType, data, tags, Nil)
+                  observableSrv.create(
+                    inputObservable.observable.copy(organisationIds = jobObs.organisationIds, relatedId = jobId),
+                    observableType,
+                    data,
+                    tags,
+                    Nil
+                  )
                 },
               inputAttachment =>
                 attachmentSrv.create(inputAttachment.name, inputAttachment.size, inputAttachment.contentType, inputAttachment.data).flatMap {
                   attachment =>
-                    observableSrv.create(inputObservable.observable, observableType, attachment, tags, Nil)
+                    observableSrv
+                      .create(
+                        inputObservable.observable.copy(organisationIds = jobObs.organisationIds, relatedId = jobId),
+                        observableType,
+                        attachment,
+                        tags,
+                        Nil
+                      )
                 }
             )
         _ = updateMetaData(richObservable.observable, inputObservable.metaData)
@@ -678,7 +700,7 @@ class Output @Inject() (
             )
         tags = inputAlert.tags.filterNot(_.isEmpty).flatMap(getTag(_).toOption).toSeq
 //        alert <- alertSrv.create(inputAlert.alert, organisation, tags, inputAlert.customFields, caseTemplate) // FIXME don't check duplicate
-        alert <- alertSrv.createEntity(inputAlert.alert)
+        alert <- alertSrv.createEntity(inputAlert.alert.copy(organisationId = organisation._id))
         _     <- alertSrv.alertOrganisationSrv.create(AlertOrganisation(), alert, organisation)
         _     <- caseTemplate.map(ct => alertSrv.alertCaseTemplateSrv.create(AlertCaseTemplate(), alert, ct)).flip
         _     <- tags.toTry(t => alertSrv.alertTagSrv.create(AlertTag(), alert, t))
@@ -691,26 +713,40 @@ class Output @Inject() (
   override def createAlertObservable(alertId: EntityId, inputObservable: InputObservable): Try[IdMapping] =
     authTransaction(inputObservable.metaData.createdBy) { implicit graph => implicit authContext =>
       logger.debug(s"Create observable ${inputObservable.dataOrAttachment.fold(identity, _.name)} in alert $alertId")
+      val tags = inputObservable.tags.filterNot(_.isEmpty).flatMap(getTag(_).toOption).toSeq
       for {
         observableType <- getObservableType(inputObservable.`type`)
-        tags = inputObservable.tags.filterNot(_.isEmpty).flatMap(getTag(_).toOption).toSeq
+        alert          <- alertSrv.getOrFail(alertId)
         richObservable <-
           inputObservable
             .dataOrAttachment
             .fold(
               dataValue =>
                 dataSrv.createEntity(Data(dataValue)).flatMap { data =>
-                  observableSrv.create(inputObservable.observable, observableType, data, tags, Nil)
+                  observableSrv.create(
+                    inputObservable.observable.copy(organisationIds = Seq(alert.organisationId), relatedId = alertId),
+                    observableType,
+                    data,
+                    tags,
+                    Nil
+                  )
                 },
               inputAttachment =>
                 attachmentSrv.create(inputAttachment.name, inputAttachment.size, inputAttachment.contentType, inputAttachment.data).flatMap {
                   attachment =>
-                    observableSrv.create(inputObservable.observable, observableType, attachment, tags, Nil)
+                    observableSrv
+                      .create(
+                        inputObservable.observable.copy(organisationIds = Seq(alert.organisationId), relatedId = alertId),
+                        observableType,
+                        attachment,
+                        tags,
+                        Nil
+                      )
                 }
             )
         _ = updateMetaData(richObservable.observable, inputObservable.metaData)
-        alert <- alertSrv.getOrFail(alertId)
-        _     <- alertSrv.alertObservableSrv.create(AlertObservable(), alert, richObservable.observable)
+
+        _ <- alertSrv.alertObservableSrv.create(AlertObservable(), alert, richObservable.observable)
       } yield IdMapping(inputObservable.metaData.id, richObservable._id)
     }
 

diff --git a/misp/connector/src/main/scala/org/thp/thehive/connector/misp/services/MispImportSrv.scala b/misp/connector/src/main/scala/org/thp/thehive/connector/misp/services/MispImportSrv.scala
@@ -23,7 +23,7 @@ import play.api.libs.json._
 
 import java.nio.file.Files
 import java.util.Date
-import javax.inject.{Inject, Named, Singleton}
+import javax.inject.{Inject, Singleton}
 import scala.concurrent.duration.DurationInt
 import scala.concurrent.{Await, ExecutionContext}
 import scala.util.{Failure, Success, Try}
@@ -91,9 +91,9 @@ class MispImportSrv @Inject() (
       .fold(observableTypeSrv.getOrFail(EntityName("other")).map(_ -> Seq.empty[String]))(Success(_))
   }
 
-  def attributeToObservable(
-      attribute: Attribute
-  )(implicit graph: Graph): List[(Observable, ObservableType with Entity, Set[String], Either[String, (String, String, Source[ByteString, _])])] =
+  def attributeToObservable(alert: Alert with Entity, attribute: Attribute)(implicit
+      graph: Graph
+  ): List[(Observable, ObservableType with Entity, Set[String], Either[String, (String, String, Source[ByteString, _])])] =
     attribute
       .`type`
       .split('|')
@@ -114,7 +114,15 @@ class MispImportSrv @Inject() (
           )
           List(
             (
-              Observable(attribute.comment, 0, ioc = false, sighted = false, ignoreSimilarity = None),
+              Observable(
+                attribute.comment,
+                0,
+                ioc = false,
+                sighted = false,
+                ignoreSimilarity = None,
+                organisationIds = Seq(alert.organisationId),
+                relatedId = alert._id
+              ),
               observableType,
               attribute.tags.map(_.name).toSet ++ additionalTags,
               Right(attribute.data.get)
@@ -126,7 +134,15 @@ class MispImportSrv @Inject() (
           )
           List(
             (
-              Observable(attribute.comment, 0, ioc = false, sighted = false, ignoreSimilarity = None),
+              Observable(
+                attribute.comment,
+                0,
+                ioc = false,
+                sighted = false,
+                ignoreSimilarity = None,
+                organisationIds = Seq(alert.organisationId),
+                relatedId = alert._id
+              ),
               observableType,
               attribute.tags.map(_.name).toSet ++ additionalTags,
               Left(attribute.value)
@@ -144,7 +160,15 @@ class MispImportSrv @Inject() (
                   s"attribute ${attribute.category}:${attribute.`type`} (${attribute.tags}) is converted to observable $observableType with tags $additionalTags"
                 )
                 (
-                  Observable(attribute.comment, 0, ioc = false, sighted = false, ignoreSimilarity = None),
+                  Observable(
+                    attribute.comment,
+                    0,
+                    ioc = false,
+                    sighted = false,
+                    ignoreSimilarity = None,
+                    organisationIds = Seq(alert.organisationId),
+                    relatedId = alert._id
+                  ),
                   observableType,
                   attribute.tags.map(_.name).toSet ++ additionalTags,
                   Left(value)
@@ -278,7 +302,7 @@ class MispImportSrv @Inject() (
     val queue =
       client
         .searchAttributes(event.id, lastSynchro)
-        .mapConcat(attributeToObservable)
+        .mapConcat(attributeToObservable(alert, _))
         .fold(
           Map.empty[
             (String, String),

diff --git a/thehive/app/org/thp/thehive/controllers/v0/AlertCtrl.scala b/thehive/app/org/thp/thehive/controllers/v0/AlertCtrl.scala
@@ -6,7 +6,7 @@ import io.scalaland.chimney.dsl._
 import javax.inject.{Inject, Named, Singleton}
 import org.thp.scalligraph.auth.AuthContext
 import org.thp.scalligraph.controllers._
-import org.thp.scalligraph.models.{Database, UMapping}
+import org.thp.scalligraph.models.{Database, Entity, UMapping}
 import org.thp.scalligraph.query._
 import org.thp.scalligraph.traversal.TraversalOps._
 import org.thp.scalligraph.traversal.{Converter, Graph, IdentityConverter, IteratorOutput, Traversal}
@@ -301,7 +301,7 @@ class AlertCtrl @Inject() (
         } yield Results.Ok(alertWithObservables.toJson)
       }
 
-  private def createObservable(observable: InputObservable)(implicit
+  private def createObservable(organisation: Organisation with Entity, observable: InputObservable)(implicit
       graph: Graph,
       authContext: AuthContext
   ): Try[Seq[RichObservable]] =
@@ -314,15 +314,20 @@ class AlertCtrl @Inject() (
               val data = Base64.getDecoder.decode(value)
               attachmentSrv
                 .create(filename, contentType, data)
-                .flatMap(attachment => observableSrv.create(observable.toObservable, attachmentType, attachment, observable.tags, Nil))
+                .flatMap(attachment =>
+                  observableSrv.create(observable.toObservable(organisation._id), attachmentType, attachment, observable.tags, Nil)
+                )
             case Array(filename, contentType) =>
               attachmentSrv
                 .create(filename, contentType, Array.emptyByteArray)
-                .flatMap(attachment => observableSrv.create(observable.toObservable, attachmentType, attachment, observable.tags, Nil))
+                .flatMap(attachment =>
+                  observableSrv.create(observable.toObservable(organisation._id), attachmentType, attachment, observable.tags, Nil)
+                )
             case data =>
               Failure(InvalidFormatAttributeError("artifacts.data", "filename;contentType;base64value", Set.empty, FString(data.mkString(";"))))
           }
-        case dataType => observable.data.toTry(d => observableSrv.create(observable.toObservable, dataType, d, observable.tags, Nil))
+        case dataType =>
+          observable.data.toTry(d => observableSrv.create(observable.toObservable(organisation._id), dataType, d, observable.tags, Nil))
       }
 }