#280 Rewrite case/alert similarity and remove substreams

TheHive-Project · Jun 14, 2018 · f45b965 · f45b965
1 parent a200bfa
commit f45b965
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 52 deletions.
diff --git a/project/Dependencies.scala b/project/Dependencies.scala
@@ -20,7 +20,7 @@ object Dependencies {
 
     val reflections = "org.reflections" % "reflections" % "0.9.11"
     val zip4j = "net.lingala.zip4j" % "zip4j" % "1.3.2"
-    val elastic4play = "org.cert-bdf" %% "elastic4play" % "1.5.0"
+    val elastic4play = "org.cert-bdf" %% "elastic4play" % "1.5.1-SNAPSHOT"
     val akkaCluster = "com.typesafe.akka" %% "akka-cluster" % "2.5.6"
     val akkaClusterTools = "com.typesafe.akka" %% "akka-cluster-tools" % "2.5.6"
   }

diff --git a/thehive-backend/app/services/AlertSrv.scala b/thehive-backend/app/services/AlertSrv.scala
@@ -38,7 +38,6 @@ object AlertSrv {
 
 @Singleton
 class AlertSrv(
-    maxSimilarCases: Int,
     templates: Map[String, String],
     alertModel: AlertModel,
     createSrv: CreateSrv,
@@ -70,7 +69,6 @@ class AlertSrv(
       connectors: ConnectorRouter,
       ec: ExecutionContext,
       mat: Materializer) = this(
-    configuration.getOptional[Int]("maxSimilarCases").getOrElse(100),
     Map.empty[String, String],
     alertModel: AlertModel,
     createSrv,
@@ -315,17 +313,17 @@ class AlertSrv(
         similarArtifacts(artifact)
           .getOrElse(Source.empty)
       }
-      .groupBy(maxSimilarCases, _.parentId)
-      .map {
-        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[(String, Int, Int)] {
-        case ((caseId, iocCount1, artifactCount1), (_, iocCount2, artifactCount2)) ⇒ (caseId, iocCount1 + iocCount2, artifactCount1 + artifactCount2)
+      .fold(Map.empty[String, (Int, Int)]) { (similarCases, artifact) ⇒
+        val caseId = artifact.parentId.getOrElse(sys.error(s"Artifact ${artifact.id} has no case !"))
+        val (iocCount, artifactCount) = similarCases.getOrElse(caseId, (0, 0))
+        if (artifact.ioc())
+          similarCases + (caseId -> ((iocCount + 1, artifactCount)))
+        else
+          similarCases + (caseId -> ((iocCount, artifactCount + 1)))
       }
-      .mergeSubstreams
+      .mapConcat(identity)
       .mapAsyncUnordered(5) {
-        case (caseId, similarIOCCount, similarArtifactCount) ⇒
+        case (caseId, (similarIOCCount, similarArtifactCount)) ⇒
           caseSrv.get(caseId).map((_, similarIOCCount, similarArtifactCount))
       }
       .filter {

diff --git a/thehive-backend/app/services/CaseSrv.scala b/thehive-backend/app/services/CaseSrv.scala
@@ -5,7 +5,7 @@ import javax.inject.{ Inject, Provider, Singleton }
 import scala.concurrent.{ ExecutionContext, Future }
 import scala.util.Try
 
-import play.api.{ Configuration, Logger }
+import play.api.Logger
 import play.api.libs.json.Json.toJsFieldJsValueWrapper
 import play.api.libs.json._
 
@@ -14,14 +14,12 @@ import akka.stream.Materializer
 import akka.stream.scaladsl.{ Sink, Source }
 import models._
 
-import org.elastic4play.InternalError
 import org.elastic4play.controllers.Fields
 import org.elastic4play.database.ModifyConfig
 import org.elastic4play.services._
 
 @Singleton
-class CaseSrv(
-    maxSimilarCases: Int,
+class CaseSrv @Inject() (
     caseModel: CaseModel,
     artifactModel: ArtifactModel,
     taskSrv: TaskSrv,
@@ -36,36 +34,6 @@ class CaseSrv(
     implicit val ec: ExecutionContext,
     implicit val mat: Materializer) {
 
-  @Inject() def this(
-      configuration: Configuration,
-      caseModel: CaseModel,
-      artifactModel: ArtifactModel,
-      taskSrv: TaskSrv,
-      auditSrv: AuditSrv,
-      alertSrvProvider: Provider[AlertSrv],
-      createSrv: CreateSrv,
-      artifactSrv: ArtifactSrv,
-      getSrv: GetSrv,
-      updateSrv: UpdateSrv,
-      deleteSrv: DeleteSrv,
-      findSrv: FindSrv,
-      ec: ExecutionContext,
-      mat: Materializer) = this(
-    configuration.getOptional[Int]("maxSimilarCases").getOrElse(100),
-    caseModel,
-    artifactModel,
-    taskSrv,
-    auditSrv,
-    alertSrvProvider,
-    createSrv,
-    artifactSrv,
-    getSrv,
-    updateSrv,
-    deleteSrv,
-    findSrv,
-    ec,
-    mat)
-
   private lazy val alertSrv = alertSrvProvider.get
   private[CaseSrv] lazy val logger = Logger(getClass)
 
@@ -184,13 +152,14 @@ class CaseSrv(
         "status" ~= "Ok"), Some("all"), Nil)
       ._1
       .flatMapConcat { artifact ⇒ artifactSrv.findSimilar(artifact, Some("all"), Nil)._1 }
-      .groupBy(maxSimilarCases, _.parentId)
-      .map { a ⇒ (a.parentId, Seq(a)) }
-      .reduce((l, r) ⇒ (l._1, r._2 ++ l._2))
-      .mergeSubstreams
+      .fold(Map.empty[String, List[Artifact]]) { (similarCases, artifact) ⇒
+        val caseId = artifact.parentId.getOrElse(sys.error(s"Artifact ${artifact.id} has no case !"))
+        val artifactList = artifact :: similarCases.getOrElse(caseId, Nil)
+        similarCases + (caseId -> artifactList)
+      }
+      .mapConcat(identity)
       .mapAsyncUnordered(5) {
-        case (Some(caseId), artifacts) ⇒ getSrv[CaseModel, Case](caseModel, caseId) map (_ → artifacts)
-        case _                         ⇒ Future.failed(InternalError("Case not found"))
+        case (caseId, artifacts) ⇒ getSrv[CaseModel, Case](caseModel, caseId) map (_ → artifacts)
       }
       .mapMaterializedValue(_ ⇒ NotUsed)
   }