#1340 Add cache in migration

TheHive-Project · Jun 13, 2020 · 28493fd · 28493fd
1 parent b663339
commit 28493fd
Show file tree

Hide file tree

Showing 4 changed files with 46 additions and 8 deletions.
diff --git a/misp/connector/src/main/scala/org/thp/thehive/connector/misp/services/MispExportSrv.scala b/misp/connector/src/main/scala/org/thp/thehive/connector/misp/services/MispExportSrv.scala
@@ -9,7 +9,7 @@ import play.api.Logger
 
 import gremlin.scala.Graph
 import javax.inject.{Inject, Singleton}
-import org.thp.misp.dto.{Attribute, Tag}
+import org.thp.misp.dto.{Attribute, Tag => MispTag}
 import org.thp.scalligraph.auth.AuthContext
 import org.thp.scalligraph.models.{Database, Entity}
 import org.thp.scalligraph.steps.StepsOps._
@@ -48,7 +48,7 @@ class MispExportSrv @Inject() (
             value = observable.data.fold(observable.attachment.get.name)(_.data),
             firstSeen = None,
             lastSeen = None,
-            tags = observable.tags.map(t => Tag(None, t.toString, Some(t.colour), None))
+            tags = observable.tags.map(t => MispTag(None, t.toString, Some(t.colour), None))
           )
       }
       .orElse {
@@ -128,7 +128,7 @@ class MispExportSrv @Inject() (
         )
       }
       org          <- organisationSrv.getOrFail(authContext.organisation)
-      createdAlert <- alertSrv.create(alert.copy(lastSyncDate = new Date(0L)), org, Set.empty, Map.empty, None)
+      createdAlert <- alertSrv.create(alert.copy(lastSyncDate = new Date(0L)), org, Seq.empty[Tag with Entity], Map.empty[String, Option[Any]], None)
       _            <- alertSrv.alertCaseSrv.create(AlertCase(), createdAlert.alert, `case`)
     } yield createdAlert
 

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
@@ -351,7 +351,7 @@ class MispImportSrv @Inject() (
           case None => // if the related alert doesn't exist, create it
             logger.debug(s"Event ${client.name}#${event.id} has no related alert for organisation ${organisation.name}")
             alertSrv
-              .create(alert, organisation, event.tags.map(_.name).toSet, Map.empty, caseTemplate)
+              .create(alert, organisation, event.tags.map(_.name).toSet, Map.empty[String, Option[Any]], caseTemplate)
               .map(_.alert)
           case Some(richAlert) =>
             logger.debug(s"Event ${client.name}#${event.id} have already been imported for organisation ${organisation.name}, updating the alert")

diff --git a/thehive/app/org/thp/thehive/services/AlertSrv.scala b/thehive/app/org/thp/thehive/services/AlertSrv.scala
@@ -54,6 +54,18 @@ class AlertSrv @Inject() (
   )(
       implicit graph: Graph,
       authContext: AuthContext
+  ): Try[RichAlert] =
+    tagNames.toTry(tagSrv.getOrCreate).flatMap(create(alert, organisation, _, customFields, caseTemplate))
+
+  def create(
+      alert: Alert,
+      organisation: Organisation with Entity,
+      tags: Seq[Tag with Entity],
+      customFields: Map[String, Option[Any]],
+      caseTemplate: Option[CaseTemplate with Entity]
+  )(
+      implicit graph: Graph,
+      authContext: AuthContext
   ): Try[RichAlert] = {
     val alertAlreadyExist = organisationSrv.get(organisation).alerts.getBySourceId(alert.`type`, alert.source, alert.sourceRef).getCount
     if (alertAlreadyExist > 0)
@@ -63,7 +75,6 @@ class AlertSrv @Inject() (
         createdAlert <- createEntity(alert)
         _            <- alertOrganisationSrv.create(AlertOrganisation(), createdAlert, organisation)
         _            <- caseTemplate.map(ct => alertCaseTemplateSrv.create(AlertCaseTemplate(), createdAlert, ct)).flip
-        tags         <- tagNames.filterNot(_.isEmpty).toTry(tagSrv.getOrCreate)
         _            <- 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))

diff --git a/thehive/app/org/thp/thehive/services/ObservableSrv.scala b/thehive/app/org/thp/thehive/services/ObservableSrv.scala
@@ -59,25 +59,52 @@ class ObservableSrv @Inject() (
   )(
       implicit graph: Graph,
       authContext: AuthContext
+  ): Try[RichObservable] =
+    tagNames.toTry(tagSrv.getOrCreate).flatMap(tags => create(observable, `type`, attachment, tags, extensions))
+
+  def create(
+      observable: Observable,
+      `type`: ObservableType with Entity,
+      attachment: Attachment with Entity,
+      tags: Seq[Tag with Entity],
+      extensions: Seq[KeyValue]
+  )(
+      implicit graph: Graph,
+      authContext: AuthContext
   ): Try[RichObservable] =
     for {
       createdObservable <- createEntity(observable)
       _                 <- observableObservableType.create(ObservableObservableType(), createdObservable, `type`)
       _                 <- observableAttachmentSrv.create(ObservableAttachment(), createdObservable, attachment)
-      tags              <- addTags(createdObservable, tagNames)
+      _                 <- tags.toTry(observableTagSrv.create(ObservableTag(), createdObservable, _))
       ext               <- addExtensions(createdObservable, extensions)
     } yield RichObservable(createdObservable, `type`, None, Some(attachment), tags, None, ext, Nil)
 
   def create(observable: Observable, `type`: ObservableType with Entity, dataValue: String, tagNames: Set[String], extensions: Seq[KeyValue])(
       implicit graph: Graph,
       authContext: AuthContext
+  ): Try[RichObservable] =
+    for {
+      tags           <- tagNames.toTry(tagSrv.getOrCreate)
+      data           <- dataSrv.create(Data(dataValue))
+      richObservable <- create(observable, `type`, data, tags, extensions)
+    } yield richObservable
+
+  def create(
+      observable: Observable,
+      `type`: ObservableType with Entity,
+      data: Data with Entity,
+      tags: Seq[Tag with Entity],
+      extensions: Seq[KeyValue]
+  )(
+      implicit graph: Graph,
+      authContext: AuthContext
   ): Try[RichObservable] =
     for {
       createdObservable <- createEntity(observable)
       _                 <- observableObservableType.create(ObservableObservableType(), createdObservable, `type`)
-      data              <- dataSrv.create(Data(dataValue))
       _                 <- observableDataSrv.create(ObservableData(), createdObservable, data)
-      tags              <- addTags(createdObservable, tagNames)
+      _                 <- tags.toTry(observableTagSrv.create(ObservableTag(), createdObservable, _))
       ext               <- addExtensions(createdObservable, extensions)
     } yield RichObservable(createdObservable, `type`, Some(data), None, tags, None, ext, Nil)