#14 Fix case merging in backend

Author: To-om

thehive-backend/app/models/Case.scala

@@ -21,12 +21,12 @@ import services.AuditedModel
 
 object CaseStatus extends Enumeration with HiveEnumeration {
   type Type = Value
-  val Ephemeral, Open, FalsePositive, TruePositive, Resolved, Deleted = Value
+  val Open, Resolved, Deleted = Value
 }
 
 object CaseResolutionStatus extends Enumeration with HiveEnumeration {
   type Type = Value
-  val Indeterminate, FalsePositive, TruePositive, Other = Value
+  val Indeterminate, FalsePositive, TruePositive, Other, Duplicated = Value
 }
 
 object CaseImpactStatus extends Enumeration with HiveEnumeration {
@@ -50,6 +50,8 @@ trait CaseAttributes { _: AttributeDef =>
   val resolutionStatus = optionalAttribute("resolutionStatus", F.enumFmt(CaseResolutionStatus), "Resolution status of the case")
   val impactStatus = optionalAttribute("impactStatus", F.enumFmt(CaseImpactStatus), "Impact status of the case")
   val summary = optionalAttribute("summary", F.textFmt, "Summary of the case, to be provided when closing a case")
+  val mergeInto = optionalAttribute("mergeInto",F.stringFmt, "Id of the case created by the merge")
+  val mergeFrom = multiAttribute("mergeFrom",F.stringFmt, "Id of the cases merged")
 }
 
 @Singleton

thehive-backend/app/models/Job.scala

@@ -23,7 +23,7 @@ trait JobAttributes { _: AttributeDef =>
   val analyzerId = attribute("analyzerId", F.stringFmt, "Analyzer", O.readonly)
   val status = attribute("status", F.enumFmt(JobStatus), "Status of the job", JobStatus.InProgress)
   val artifactId = attribute("artifactId", F.stringFmt, "Original artifact on which this job was executed", O.readonly)
-  val startDate = attribute("startDate", F.dateFmt, "Timestamp of the job start", O.model)
+  val startDate = attribute("startDate", F.dateFmt, "Timestamp of the job start") // , O.model)
   val endDate = optionalAttribute("endDate", F.dateFmt, "Timestamp of the job completion (or fail)")
   val report = optionalAttribute("report", F.textFmt, "Analysis result", O.unaudited)
 

thehive-backend/app/services/CaseMergeSrv.scala

@@ -22,6 +22,11 @@ import org.elastic4play.services.JsonFormat.log
 import org.elastic4play.services.QueryDSL
 
 import models.{ Artifact, ArtifactStatus, Case, CaseImpactStatus, CaseResolutionStatus, CaseStatus, JobStatus, Task }
+import play.api.Logger
+import scala.util.Success
+import scala.util.Failure
+import models.TaskStatus
+import models.LogStatus
 
 @Singleton
 class CaseMergeSrv @Inject() (caseSrv: CaseSrv,
@@ -32,9 +37,20 @@ class CaseMergeSrv @Inject() (caseSrv: CaseSrv,
                               implicit val ec: ExecutionContext,
                               implicit val mat: Materializer) {
 
+  lazy val logger = Logger(getClass)
+
   import QueryDSL._
-  private[services] def concat[E](entities: Seq[E], sep: String, getId: E => Long, getStr: E => String) = {
-    JsString(entities.map(e => s"#${getId(e)}:${getStr(e)}").mkString(sep))
+  private[services] def concat[E](entities: Seq[E], sep: String, getId: E ⇒ Long, getStr: E ⇒ String) = {
+    JsString(entities.map(e ⇒ s"#${getId(e)}:${getStr(e)}").mkString(sep))
+  }
+
+  private[services] def concatCaseDescription(cases: Seq[Case]) = {
+    val str = cases
+      .map { caze ⇒
+        s"#### ${caze.title()} ([#${caze.caseId()}](#/case/${caze.id}/details))\n\n${caze.description()}"
+      }
+      .mkString("\n  \n")
+    JsString(str)
   }
 
   private[services] def firstDate(dates: Seq[Date]) = Json.toJson(dates.min)
@@ -43,39 +59,39 @@ class CaseMergeSrv @Inject() (caseSrv: CaseSrv,
     val resolutionStatus = cases
       .map(_.resolutionStatus())
       .reduce[Option[CaseResolutionStatus.Type]] {
-        case (None, s)                                     => s
-        case (s, None)                                     => s
-        case (Some(CaseResolutionStatus.Other), s)         => s
-        case (s, Some(CaseResolutionStatus.Other))         => s
-        case (Some(CaseResolutionStatus.FalsePositive), s) => s
-        case (s, Some(CaseResolutionStatus.FalsePositive)) => s
-        case (Some(CaseResolutionStatus.Indeterminate), s) => s
-        case (s, Some(CaseResolutionStatus.Indeterminate)) => s
-        case (s, _)                                        => s //TruePositive
+        case (None, s)                                     ⇒ s
+        case (s, None)                                     ⇒ s
+        case (Some(CaseResolutionStatus.Other), s)         ⇒ s
+        case (s, Some(CaseResolutionStatus.Other))         ⇒ s
+        case (Some(CaseResolutionStatus.FalsePositive), s) ⇒ s
+        case (s, Some(CaseResolutionStatus.FalsePositive)) ⇒ s
+        case (Some(CaseResolutionStatus.Indeterminate), s) ⇒ s
+        case (s, Some(CaseResolutionStatus.Indeterminate)) ⇒ s
+        case (s, _)                                        ⇒ s //TruePositive
       }
-    resolutionStatus.map(s => JsString(s.toString))
+    resolutionStatus.map(s ⇒ JsString(s.toString))
   }
 
   private[services] def mergeImpactStatus(cases: Seq[Case]) = {
     val impactStatus = cases
       .map(_.impactStatus())
       .reduce[Option[CaseImpactStatus.Type]] {
-        case (None, s)                                 => s
-        case (s, None)                                 => s
-        case (Some(CaseImpactStatus.NotApplicable), s) => s
-        case (s, Some(CaseImpactStatus.NotApplicable)) => s
-        case (Some(CaseImpactStatus.NoImpact), s)      => s
-        case (s, Some(CaseImpactStatus.NoImpact))      => s
-        case (s, _)                                    => s // WithImpact
+        case (None, s)                                 ⇒ s
+        case (s, None)                                 ⇒ s
+        case (Some(CaseImpactStatus.NotApplicable), s) ⇒ s
+        case (s, Some(CaseImpactStatus.NotApplicable)) ⇒ s
+        case (Some(CaseImpactStatus.NoImpact), s)      ⇒ s
+        case (s, Some(CaseImpactStatus.NoImpact))      ⇒ s
+        case (s, _)                                    ⇒ s // WithImpact
       }
-    impactStatus.map(s => JsString(s.toString))
+    impactStatus.map(s ⇒ JsString(s.toString))
   }
 
   private[services] def mergeSummary(cases: Seq[Case]) = {
     val summary = cases
-      .flatMap(c => c.summary().map(_ -> c.caseId()))
+      .flatMap(c ⇒ c.summary().map(_ -> c.caseId()))
       .map {
-        case (summary, caseId) => s"#$caseId:$summary"
+        case (summary, caseId) ⇒ s"#$caseId:$summary"
       }
     if (summary.isEmpty)
       None
@@ -85,13 +101,13 @@ class CaseMergeSrv @Inject() (caseSrv: CaseSrv,
 
   private[services] def mergeMetrics(cases: Seq[Case]): JsObject = {
     val metrics = for {
-      caze <- cases
-      metrics <- caze.metrics()
-      metricsObject <- metrics.asOpt[JsObject]
+      caze ← cases
+      metrics ← caze.metrics()
+      metricsObject ← metrics.asOpt[JsObject]
     } yield metricsObject
 
-    val mergedMetrics: Seq[(String, JsValue)] = metrics.flatMap(_.keys).distinct.map { key =>
-      val metricValues = metrics.flatMap(m => (m \ key).asOpt[BigDecimal])
+    val mergedMetrics: Seq[(String, JsValue)] = metrics.flatMap(_.keys).distinct.map { key ⇒
+      val metricValues = metrics.flatMap(m ⇒ (m \ key).asOpt[BigDecimal])
       if (metricValues.size != 1)
         key -> JsNull
       else
@@ -105,24 +121,30 @@ class CaseMergeSrv @Inject() (caseSrv: CaseSrv,
 
   private[services] def mergeLogs(oldTask: Task, newTask: Task)(implicit authContext: AuthContext): Future[Done] = {
     logSrv.find("_parent" ~= oldTask.id, Some("all"), Nil)._1
-      .mapAsyncUnordered(5) { log =>
+      .mapAsyncUnordered(5) { log ⇒
         logSrv.create(newTask, baseFields(log))
       }
       .runWith(Sink.ignore)
   }
 
   private[services] def mergeTasksAndLogs(newCase: Case, cases: Seq[Case])(implicit authContext: AuthContext): Future[Done] = {
-    taskSrv.find(or(cases.map("_parent" ~= _.id)), Some("all"), Nil)._1
-      .mapAsyncUnordered(5) { task =>
-        taskSrv.create(newCase, baseFields(task)).map(task -> _)
-      }
+    val (tasks, futureTaskCount) = taskSrv.find(and(parent("case", withId(cases.map(_.id): _*)), "status" ~!= TaskStatus.Cancel), Some("all"), Nil)
+    futureTaskCount.foreach(count ⇒ logger.info(s"Creating $count task(s):"))
+    tasks
+      .mapAsyncUnordered(5) { task ⇒ taskSrv.create(newCase, baseFields(task)).map(task -> _) }
       .flatMapConcat {
-        case (oldTask, newTask) =>
-          logSrv.find("_parent" ~= oldTask.id, Some("all"), Nil)._1
-            .map(_ -> newTask)
+        case (oldTask, newTask) ⇒
+          logger.info(s"\ttask : ${oldTask.id} -> ${newTask.id} : ${newTask.title()}")
+          val (logs, futureLogCount) = logSrv.find(and(parent("case_task", withId(oldTask.id)), "status" ~!= LogStatus.Deleted), Some("all"), Nil)
+          futureLogCount.foreach { count ⇒ logger.info(s"Creating $count log(s) in task ${newTask.id}") }
+          logs.map(_ -> newTask)
       }
       .mapAsyncUnordered(5) {
-        case (log, task) => logSrv.create(task, baseFields(log))
+        case (log, task) ⇒
+          val fields = log.attachment().fold(baseFields(log)) { a ⇒
+            baseFields(log).set("attachment", AttachmentInputValue(a.name, a.hashes, a.size, a.contentType, a.id))
+          }
+          logSrv.create(task, fields)
       }
       .runWith(Sink.ignore)
   }
@@ -131,71 +153,86 @@ class CaseMergeSrv @Inject() (caseSrv: CaseSrv,
     val status = artifacts
       .map(_.status())
       .reduce[ArtifactStatus.Type] {
-        case (ArtifactStatus.Deleted, s) => s
-        case (s, _)                      => s
+        case (ArtifactStatus.Deleted, s) ⇒ s
+        case (s, _)                      ⇒ s
       }
       .toString
     JsString(status)
   }
 
   private[services] def mergeJobs(newArtifact: Artifact, artifacts: Seq[Artifact])(implicit authContext: AuthContext): Future[Done] = {
-    jobSrv.find(and(or(artifacts.map("_parent" ~= _.id)), "status" ~= JobStatus.Success), Some("all"), Nil)._1
-      .mapAsyncUnordered(5) { job =>
+    jobSrv.find(and(parent("case_artifact", withId(artifacts.map(_.id): _*)), "status" ~= JobStatus.Success), Some("all"), Nil)._1
+      .mapAsyncUnordered(5) { job ⇒
         jobSrv.create(newArtifact, baseFields(job))
       }
       .runWith(Sink.ignore)
   }
 
   private[services] def mergeArtifactsAndJobs(newCase: Case, cases: Seq[Case])(implicit authContext: AuthContext): Future[Done] = {
-    val caseMap = cases.map(c => c.id -> c).toMap
-    val caseFilter = or(cases.map("_parent" ~= _.id))
+    val caseMap = cases.map(c ⇒ c.id -> c).toMap
+    val caseFilter = and(parent("case", withId(cases.map(_.id): _*)), "status" ~= "Ok")
     // Find artifacts hold by cases
-    artifactSrv.find(caseFilter, Some("all"), Nil)._1
-      .map { artifact =>
+    val (artifacts, futureArtifactCount) = artifactSrv.find(caseFilter, Some("all"), Nil)
+    futureArtifactCount.foreach { count ⇒ log.info(s"Found $count artifact(s) in merging cases") }
+    artifacts
+      .mapAsyncUnordered(5) { artifact ⇒
         // For each artifact find similar artifacts
         val dataFilter = artifact.data().map("data" ~= _) orElse artifact.attachment().map("attachment.id" ~= _.id)
         val filter = and(caseFilter,
           "status" ~= "Ok",
           "dataType" ~= artifact.dataType(),
           dataFilter.get)
-        artifactSrv.find(filter, Some("all"), Nil)._1
-          .runWith(Sink.seq)
-          .flatMap { sameArtifacts =>
-            // Same artifacts are merged
-            val firstArtifact = sameArtifacts.head
-            val fields = firstArtifact.attachment().fold(Fields.empty) { a =>
-              Fields.empty.set("attachment", AttachmentInputValue(a.name, a.hashes, a.size, a.contentType, a.id))
-            }
-              .set("data", firstArtifact.data().map(JsString))
-              .set("dataType", firstArtifact.dataType())
-              .set("message", concat[Artifact](sameArtifacts, "\n  \n", a => caseMap(a.parentId.get).caseId(), _.message()))
-              .set("startDate", firstDate(sameArtifacts.map(_.startDate())))
-              .set("tlp", JsNumber(sameArtifacts.map(_.tlp()).min))
-              .set("tags", JsArray(sameArtifacts.flatMap(_.tags()).map(JsString)))
-              .set("ioc", JsBoolean(sameArtifacts.map(_.ioc()).reduce(_ || _)))
-              .set("status", mergeArtifactStatus(sameArtifacts))
-            // Merged artifact is created under new case
-            artifactSrv
-              .create(newCase, fields)
-              // Then jobs are imported
-              .flatMap { newArtifact =>
-                mergeJobs(newArtifact, sameArtifacts)
-              }
-              // Errors are logged and ignored (probably document already exists)
-              .recover {
-                case error =>
-                  log.warn("Artifact creation fail", error)
-                  Done
-              }
+
+        val (artifacts, futureArtifactCount) = artifactSrv.find(filter, Some("all"), Nil)
+        futureArtifactCount.foreach { count ⇒
+          logger.debug(s"${count} identical artifact(s) found (${artifact.dataType()}):${(artifact.data() orElse artifact.attachment().map(_.name)).get}")
+        }
+        artifacts.runWith(Sink.seq)
+      }
+      .mapAsync(5) { sameArtifacts ⇒
+        // Same artifacts are merged
+        val firstArtifact = sameArtifacts.head
+        val fields = firstArtifact.attachment().fold(Fields.empty) { a ⇒
+          Fields.empty.set("attachment", AttachmentInputValue(a.name, a.hashes, a.size, a.contentType, a.id))
+        }
+          .set("data", firstArtifact.data().map(JsString))
+          .set("dataType", firstArtifact.dataType())
+          .set("message", concat[Artifact](sameArtifacts, "\n  \n", a ⇒ caseMap(a.parentId.get).caseId(), _.message()))
+          .set("startDate", firstDate(sameArtifacts.map(_.startDate())))
+          .set("tlp", JsNumber(sameArtifacts.map(_.tlp()).min))
+          .set("tags", JsArray(sameArtifacts.flatMap(_.tags()).map(JsString)))
+          .set("ioc", JsBoolean(sameArtifacts.map(_.ioc()).reduce(_ || _)))
+          .set("status", mergeArtifactStatus(sameArtifacts))
+        // Merged artifact is created under new case
+        artifactSrv
+          .create(newCase, fields)
+          .map(a ⇒ List(a -> sameArtifacts))
+          // Errors are logged and ignored (probably document already exists)
+          .recover {
+            case e ⇒
+              logger.warn("Artifact creation fail", e)
+              Nil
           }
       }
+      .mapConcat(identity)
+      .mapAsyncUnordered(5) {
+        case (newArtifact, sameArtifacts) ⇒
+          // Then jobs are imported
+          mergeJobs(newArtifact, sameArtifacts)
+            .recover {
+              case error ⇒
+                logger.error("Log creation fail", error)
+                Done
+            }
+      }
       .runWith(Sink.ignore)
   }
 
   private[services] def mergeCases(cases: Seq[Case])(implicit authContext: AuthContext): Future[Case] = {
+    logger.info("Merging cases: " + cases.map(c ⇒ s"#${c.caseId()}:${c.title()}").mkString(" / "))
     val fields = Fields.empty
       .set("title", concat[Case](cases, " / ", _.caseId(), _.title()))
-      .set("description", concat[Case](cases, "\n  \n", _.caseId(), _.description()))
+      .set("description", concatCaseDescription(cases))
       .set("severity", JsNumber(cases.map(_.severity()).max))
       .set("startDate", firstDate(cases.map(_.startDate())))
       .set("tags", JsArray(cases.flatMap(_.tags()).distinct.map(JsString)))
@@ -207,15 +244,30 @@ class CaseMergeSrv @Inject() (caseSrv: CaseSrv,
       .set("resolutionStatus", mergeResolutionStatus(cases))
       .set("impactStatus", mergeImpactStatus(cases))
       .set("summary", mergeSummary(cases))
+      .set("mergeFrom", JsArray(cases.map(c ⇒ JsString(c.id))))
     caseSrv.create(fields)
   }
 
+  def markCaseAsDuplicated(caseIds: Seq[String], mergeCaseId: String)(implicit authContext: AuthContext): Future[Unit] = {
+    caseSrv.bulkUpdate(caseIds, Fields.empty
+      .set("mergeInto", mergeCaseId)
+      .set("status", CaseStatus.Resolved.toString)
+      .set("resolutionStatus", CaseResolutionStatus.Duplicated.toString))
+      .map(_.foreach {
+        case Success(_) ⇒ Done
+        case Failure(error) ⇒
+          log.error("Case update fail", error)
+          Done
+      })
+  }
+
   def merge(caseIds: String*)(implicit authContext: AuthContext): Future[Case] = {
     for {
-      cases <- Future.sequence(caseIds.map(caseSrv.get))
-      newCase <- mergeCases(cases)
-      _ <- mergeTasksAndLogs(newCase, cases)
-      _ <- mergeArtifactsAndJobs(newCase, cases)
+      cases ← Future.sequence(caseIds.map(caseSrv.get))
+      newCase ← mergeCases(cases)
+      _ ← mergeTasksAndLogs(newCase, cases)
+      _ ← mergeArtifactsAndJobs(newCase, cases)
+      _ ← markCaseAsDuplicated(caseIds, newCase.id)
     } yield newCase
   }
 }