#1676 Fix creation of duplicated audit during migration

TheHive-Project · Nov 20, 2020 · c89cdf2 · c89cdf2
1 parent 14281be
commit c89cdf2
Show file tree

Hide file tree

Showing 5 changed files with 58 additions and 19 deletions.
diff --git a/ScalliGraph b/ScalliGraph
diff --git a/migration/src/main/scala/org/thp/thehive/migration/Input.scala b/migration/src/main/scala/org/thp/thehive/migration/Input.scala
@@ -137,9 +137,11 @@ trait Input {
   def listAction(filter: Filter): Source[Try[(String, InputAction)], NotUsed]
   def countAction(filter: Filter): Future[Long]
   def listAction(entityId: String): Source[Try[(String, InputAction)], NotUsed]
+  def listActions(entityIds: Seq[String]): Source[Try[(String, InputAction)], NotUsed]
   def countAction(entityId: String): Future[Long]
   def listAudit(filter: Filter): Source[Try[(String, InputAudit)], NotUsed]
   def countAudit(filter: Filter): Future[Long]
   def listAudit(entityId: String, filter: Filter): Source[Try[(String, InputAudit)], NotUsed]
+  def listAudits(entityIds: Seq[String], filter: Filter): Source[Try[(String, InputAudit)], NotUsed]
   def countAudit(entityId: String, filter: Filter): Future[Long]
 }
diff --git a/migration/src/main/scala/org/thp/thehive/migration/MigrationOps.scala b/migration/src/main/scala/org/thp/thehive/migration/MigrationOps.scala
@@ -7,7 +7,7 @@ import org.thp.scalligraph.{EntityId, NotFoundError, RichOptionTry}
 import org.thp.thehive.migration.dto.{InputAlert, InputAudit, InputCase, InputCaseTemplate}
 import play.api.Logger
 
-import scala.collection.{immutable, mutable}
+import scala.collection.mutable
 import scala.concurrent.{ExecutionContext, Future}
 import scala.util.{Failure, Success, Try}
 
@@ -251,10 +251,10 @@ trait MigrationOps {
               output.createJobObservable
             )
             caseEntitiesIds = caseTaskIds ++ caseTaskLogIds ++ caseObservableIds ++ jobIds ++ jobObservableIds :+ caseId
-            actionSource    = Source(caseEntitiesIds.to[immutable.Iterable]).flatMapConcat(id => input.listAction(id.inputId))
+            actionSource    = input.listActions(caseEntitiesIds.map(_.inputId).distinct)
             actionIds <- migrateWithParent("Action", caseEntitiesIds, actionSource, output.createAction)
             caseEntitiesAuditIds = caseEntitiesIds ++ actionIds
-            auditSource          = Source(caseEntitiesAuditIds.to[immutable.Iterable]).flatMapConcat(id => input.listAudit(id.inputId, filter))
+            auditSource          = input.listAudits(caseEntitiesAuditIds.map(_.inputId).distinct, filter)
             _ <- migrateAudit(caseEntitiesAuditIds, auditSource, output.createAudit)
           } yield Some(caseId)
       }
@@ -282,10 +282,10 @@ trait MigrationOps {
               output.createAlertObservable
             )
             alertEntitiesIds = alertId +: alertObservableIds
-            actionSource     = Source(alertEntitiesIds.to[immutable.Iterable]).flatMapConcat(id => input.listAction(id.inputId))
+            actionSource     = input.listActions(alertEntitiesIds.map(_.inputId).distinct)
             actionIds <- migrateWithParent("Action", alertEntitiesIds, actionSource, output.createAction)
             alertEntitiesAuditIds = alertEntitiesIds ++ actionIds
-            auditSource           = Source(alertEntitiesAuditIds.to[immutable.Iterable]).flatMapConcat(id => input.listAudit(id.inputId, filter))
+            auditSource           = input.listAudits(alertEntitiesAuditIds.map(_.inputId).distinct, filter)
             _ <- migrateAudit(alertEntitiesAuditIds, auditSource, output.createAudit)
           } yield ()
       }

diff --git a/migration/src/main/scala/org/thp/thehive/migration/th3/Input.scala b/migration/src/main/scala/org/thp/thehive/migration/th3/Input.scala
@@ -639,7 +639,16 @@ class Input @Inject() (configuration: Configuration, dbFind: DBFind, dbGet: DBGe
     dbFind(Some("0-0"), Nil)(indexName => search(indexName).query(termQuery("relations", "action")))._2
 
   override def listAction(entityId: String): Source[Try[(String, InputAction)], NotUsed] =
-    dbFind(Some("all"), Nil)(indexName => search(indexName).query(bool(Seq(termQuery("relations", "action"), idsQuery(entityId)), Nil, Nil)))
+    dbFind(Some("all"), Nil)(indexName =>
+      search(indexName).query(bool(Seq(termQuery("relations", "action"), termQuery("objectId", entityId)), Nil, Nil))
+    )
+      ._1
+      .read[(String, InputAction)]
+
+  override def listActions(entityIds: Seq[String]): Source[Try[(String, InputAction)], NotUsed] =
+    dbFind(Some("all"), Nil)(indexName =>
+      search(indexName).query(bool(Seq(termQuery("relations", "action"), termsQuery("objectId", entityIds)), Nil, Nil))
+    )
       ._1
       .read[(String, InputAction)]
 
@@ -679,11 +688,34 @@ class Input @Inject() (configuration: Configuration, dbFind: DBFind, dbGet: DBGe
 
   override def listAudit(entityId: String, filter: Filter): Source[Try[(String, InputAudit)], NotUsed] =
     dbFind(Some("all"), Nil)(indexName =>
-      search(indexName).query(bool(auditFilter(filter) :+ termQuery("relations", "audit") :+ termQuery("objectId", entityId), Nil, Nil))
+      search(indexName).query(
+        bool(
+          auditFilter(filter) ++ auditIncludeFilter(filter) :+ termQuery("relations", "audit") :+ termQuery("objectId", entityId),
+          Nil,
+          auditExcludeFilter(filter)
+        )
+      )
+    )._1.read[(String, InputAudit)]
+
+  override def listAudits(entityIds: Seq[String], filter: Filter): Source[Try[(String, InputAudit)], NotUsed] =
+    dbFind(Some("all"), Nil)(indexName =>
+      search(indexName).query(
+        bool(
+          auditFilter(filter) ++ auditIncludeFilter(filter) :+ termQuery("relations", "audit") :+ termsQuery("objectId", entityIds),
+          Nil,
+          auditExcludeFilter(filter)
+        )
+      )
     )._1.read[(String, InputAudit)]
 
   def countAudit(entityId: String, filter: Filter): Future[Long] =
     dbFind(Some("0-0"), Nil)(indexName =>
-      search(indexName).query(bool(auditFilter(filter) :+ termQuery("relations", "audit") :+ termQuery("objectId", entityId), Nil, Nil))
+      search(indexName).query(
+        bool(
+          auditFilter(filter) ++ auditIncludeFilter(filter) :+ termQuery("relations", "audit") :+ termQuery("objectId", entityId),
+          Nil,
+          auditExcludeFilter(filter)
+        )
+      )
     )._2
 }
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
@@ -32,6 +32,7 @@ import play.api.{Configuration, Environment, Logger}
 import scala.collection.JavaConverters._
 import scala.concurrent.ExecutionContext
 import scala.util.{Failure, Success, Try}
+import org.thp.thehive.controllers.v1.Conversion._
 
 object Output {
 
@@ -220,15 +221,15 @@ class Output @Inject() (
       alerts.nonEmpty
     )
       logger.info(s"""Already migrated:
-                     | ${profiles.size} profiles\n
-                     | ${organisations.size} organisations\n
-                     | ${users.size} users\n
-                     | ${impactStatuses.size} impactStatuses\n
-                     | ${resolutionStatuses.size} resolutionStatuses\n
-                     | ${observableTypes.size} observableTypes\n
-                     | ${customFields.size} customFields\n
-                     | ${caseTemplates.size} caseTemplates\n
-                     | ${caseNumbers.size} caseNumbers\n
+                     | ${profiles.size} profiles
+                     | ${organisations.size} organisations
+                     | ${users.size} users
+                     | ${impactStatuses.size} impactStatuses
+                     | ${resolutionStatuses.size} resolutionStatuses
+                     | ${observableTypes.size} observableTypes
+                     | ${customFields.size} customFields
+                     | ${caseTemplates.size} caseTemplates
+                     | ${caseNumbers.size} caseNumbers
                      | ${alerts.size} alerts""".stripMargin)
   }
 
@@ -579,7 +580,9 @@ class Output @Inject() (
       for {
         task <- taskSrv.getOrFail(taskId)
         _ = logger.debug(s"Create log in task ${task.title}")
-        log <- logSrv.create(inputLog.log, task)
+        log <- logSrv.createEntity(inputLog.log)
+        _   <- logSrv.taskLogSrv.create(TaskLog(), task, log)
+        _   <- auditSrv.log.create(log, task, RichLog(log, Nil).toJson)
         _ = updateMetaData(log, inputLog.metaData)
         _ <- inputLog.attachments.toTry { inputAttachment =>
           attachmentSrv.create(inputAttachment.name, inputAttachment.size, inputAttachment.contentType, inputAttachment.data).flatMap { attachment =>
@@ -717,6 +720,7 @@ class Output @Inject() (
       case "Log"        => logSrv.getOrFail(entityId)
       case "Alert"      => alertSrv.getOrFail(entityId)
       case "Job"        => jobSrv.getOrFail(entityId)
+      case "Action"     => actionSrv.getOrFail(entityId)
       case _            => Failure(BadRequestError(s"objectType $entityType is not recognised"))
     }
 
@@ -744,6 +748,7 @@ class Output @Inject() (
           case "Alert"                                        => "Alert"
           case "Log" | "Task" | "Observable" | "Case" | "Job" => "Case"
           case "User"                                         => "User"
+          case "Action"                                       => "Action" // FIXME
           case other =>
             logger.error(s"Unknown object type: $other")
             other
+20 −9		core/src/main/scala/org/thp/scalligraph/traversal/TraversalOps.scala
+1 −1		project/build.properties
+1 −1		project/plugins.sbt