#272 Exclude merged cases in alert similarity

TheHive-Project · Jul 27, 2017 · f2c296b · f2c296b
1 parent b18ea3e
commit f2c296b
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 18 deletions.
diff --git a/thehive-backend/app/controllers/AlertCtrl.scala b/thehive-backend/app/controllers/AlertCtrl.scala
@@ -13,6 +13,7 @@ import play.api.http.Status
 import play.api.libs.json.{ JsArray, JsObject, Json }
 import play.api.mvc.{ Action, AnyContent, Controller }
 import services.{ AlertSrv, CaseSrv }
+import services.JsonFormat.caseSimilarityWrites
 
 import scala.concurrent.{ ExecutionContext, Future }
 import scala.util.Try

diff --git a/thehive-backend/app/services/AlertSrv.scala b/thehive-backend/app/services/AlertSrv.scala
@@ -11,6 +11,7 @@ import models._
 import org.elastic4play.InternalError
 import org.elastic4play.controllers.{ Fields, FileInputValue }
 import org.elastic4play.services.JsonFormat.attachmentFormat
+import org.elastic4play.services.QueryDSL.{ groupByField, parent, selectCount, withId }
 import org.elastic4play.services._
 import play.api.libs.json._
 import play.api.{ Configuration, Logger }
@@ -280,35 +281,34 @@ class AlertSrv(
       } yield artifactSrv.findSimilar(dataType, data, None, Some("all"), Nil)._1
     }
 
-    def getCaseAndArtifactCount(caseId: String): Future[(Case, Int, Int)] = {
-      import org.elastic4play.services.QueryDSL._
-      for {
-        caze ← caseSrv.get(caseId)
-        artifactCountJs ← artifactSrv.stats(parent("case", withId(caseId)), Seq(groupByField("ioc", selectCount)))
-        iocCount = (artifactCountJs \ "1" \ "count").asOpt[Int].getOrElse(0)
-        artifactCount = (artifactCountJs \\ "count").map(_.as[Int]).sum
-      } yield (caze, iocCount, artifactCount)
-    }
-
     Source(alert.artifacts().to[immutable.Iterable])
       .flatMapConcat { artifact ⇒
         similarArtifacts(artifact)
           .getOrElse(Source.empty)
       }
       .groupBy(100, _.parentId)
       .map {
-        case a if a.ioc() ⇒ (a.parentId, 1, 1)
-        case a            ⇒ (a.parentId, 0, 1)
+        case a if a.ioc() ⇒ (a.parentId.getOrElse(sys.error("Artifact without case !")), 1, 1)
+        case a            ⇒ (a.parentId.getOrElse(sys.error("Artifact without case !")), 0, 1)
       }
-      .reduce[(Option[String], Int, Int)] {
-        case ((caze, iocCount1, artifactCount1), (_, iocCount2, artifactCount2)) ⇒ (caze, iocCount1 + iocCount2, artifactCount1 + artifactCount2)
+      .reduce[(String, Int, Int)] {
+        case ((caseId, iocCount1, artifactCount1), (_, iocCount2, artifactCount2)) ⇒ (caseId, iocCount1 + iocCount2, artifactCount1 + artifactCount2)
       }
       .mergeSubstreams
       .mapAsyncUnordered(5) {
-        case (Some(caseId), similarIOCCount, similarArtifactCount) ⇒
-          getCaseAndArtifactCount(caseId).map {
-            case (caze, iocCount, artifactCount) ⇒ CaseSimilarity(caze, similarIOCCount, iocCount, similarArtifactCount, artifactCount)
-          }
+        case (caseId, similarIOCCount, similarArtifactCount) ⇒
+          caseSrv.get(caseId).map((_, similarIOCCount, similarArtifactCount))
+      }
+      .filter {
+        case (caze, _, _) ⇒ caze.status() != CaseStatus.Deleted && caze.resolutionStatus != CaseResolutionStatus.Duplicated
+      }
+      .mapAsyncUnordered(5) {
+        case (caze, similarIOCCount, similarArtifactCount) ⇒
+          for {
+            artifactCountJs ← artifactSrv.stats(parent("case", withId(caze.id)), Seq(groupByField("ioc", selectCount)))
+            iocCount = (artifactCountJs \ "1" \ "count").asOpt[Int].getOrElse(0)
+            artifactCount = (artifactCountJs \\ "count").map(_.as[Int]).sum
+          } yield CaseSimilarity(caze, similarIOCCount, iocCount, similarArtifactCount, artifactCount)
         case _ ⇒ Future.failed(InternalError("Case not found"))
       }
       .runWith(Sink.seq)