#232 Add event similarity

project/Dependencies.scala

@@ -30,6 +30,6 @@ object Dependencies {
     val reflections = "org.reflections" % "reflections" % "0.9.10"
     val zip4j = "net.lingala.zip4j" % "zip4j" % "1.3.2"
     val akkaTest = "com.typesafe.akka" %% "akka-stream-testkit" % "2.4.4"
-    val elastic4play = "org.cert-bdf" %% "elastic4play" % "1.1.5"
+    val elastic4play = "org.cert-bdf" %% "elastic4play" % "1.1.6-SNAPSHOT"
   }
 }

thehive-backend/app/controllers/AlertCtrl.scala

@@ -10,9 +10,10 @@ import org.elastic4play.services._
 import org.elastic4play.{ BadRequestError, Timed }
 import play.api.Logger
 import play.api.http.Status
-import play.api.libs.json.JsArray
+import play.api.libs.json.{ JsArray, JsObject, Json }
 import play.api.mvc.{ Action, AnyContent, Controller }
 import services.AlertSrv
+import services.JsonFormat.caseSimilarityWrites
 
 import scala.concurrent.{ ExecutionContext, Future }
 import scala.util.Try
@@ -37,15 +38,30 @@ class AlertCtrl @Inject() (
 
   @Timed
   def get(id: String): Action[AnyContent] = authenticated(Role.read).async { implicit request ⇒
-    val withStats = for {
-      statsValues ← request.queryString.get("nstats")
-      firstValue ← statsValues.headOption
-    } yield Try(firstValue.toBoolean).getOrElse(firstValue == "1")
+    val withStats = request
+      .queryString
+      .get("nstats")
+      .flatMap(_.headOption)
+      .exists(v ⇒ Try(v.toBoolean).getOrElse(v == "1"))
+
+    val withSimilarity = request
+      .queryString
+      .get("similarity")
+      .flatMap(_.headOption)
+      .exists(v ⇒ Try(v.toBoolean).getOrElse(v == "1"))
+    println(s"similarity=$withSimilarity")
 
     for {
       alert ← alertSrv.get(id)
-      alertsWithStats ← auxSrv.apply(alert, 0, withStats.getOrElse(false), removeUnaudited = false)
-    } yield renderer.toOutput(OK, alertsWithStats)
+      alertsWithStats ← auxSrv.apply(alert, 0, withStats, removeUnaudited = false)
+      similarCases ← if (withSimilarity)
+        alertSrv.similarCases(alert)
+          .map(sc ⇒ Json.obj("similarCases" → Json.toJson(sc)))
+      else Future.successful(JsObject(Nil))
+    } yield {
+      println(s"Similar cases = $similarCases")
+      renderer.toOutput(OK, alertsWithStats ++ similarCases)
+    }
   }
 
   @Timed

thehive-backend/app/services/AlertSrv.scala

@@ -8,18 +8,23 @@ import akka.stream.Materializer
 import akka.stream.scaladsl.{ Sink, Source }
 import connectors.ConnectorRouter
 import models._
+import org.elastic4play.InternalError
 import org.elastic4play.controllers.{ Fields, FileInputValue }
 import org.elastic4play.services._
-import play.api.{ Configuration, Logger }
+import org.elastic4play.utils.Hasher
 import play.api.libs.json._
+import play.api.{ Configuration, Logger }
 
+import scala.collection.immutable
 import scala.concurrent.{ ExecutionContext, Future }
 import scala.util.{ Failure, Try }
 
 trait AlertTransformer {
   def createCase(alert: Alert)(implicit authContext: AuthContext): Future[Case]
 }
 
+case class CaseSimilarity(caze: Case, similarIOCCount: Int, iocCount: Int, similarArtifactCount: Int, artifactCount: Int)
+
 class AlertSrv(
     templates: Map[String, String],
     alertModel: AlertModel,
@@ -33,6 +38,7 @@ class AlertSrv(
     caseTemplateSrv: CaseTemplateSrv,
     attachmentSrv: AttachmentSrv,
     connectors: ConnectorRouter,
+    hashAlg: Seq[String],
     implicit val ec: ExecutionContext,
     implicit val mat: Materializer) extends AlertTransformer {
 
@@ -63,6 +69,7 @@ class AlertSrv(
     caseTemplateSrv,
     attachmentSrv,
     connectors,
+    (configuration.getString("datastore.hash.main").get +: configuration.getStringSeq("datastore.hash.extra").get).distinct,
     ec,
     mat)
 
@@ -111,7 +118,7 @@ class AlertSrv(
     }
   }
 
-  def getCaseTemplate(alert: Alert) = {
+  def getCaseTemplate(alert: Alert): Future[Option[CaseTemplate]] = {
     val templateName = alert.caseTemplate()
       .orElse(templates.get(alert.tpe()))
       .getOrElse(alert.tpe())
@@ -209,4 +216,57 @@ class AlertSrv(
   def setFollowAlert(alertId: String, follow: Boolean)(implicit authContext: AuthContext): Future[Alert] = {
     updateSrv[AlertModel, Alert](alertModel, alertId, Fields(Json.obj("follow" → follow)))
   }
+
+  def similarCases(alert: Alert): Future[Seq[CaseSimilarity]] = {
+    def similarArtifacts(artifact: JsObject): Option[Source[Artifact, NotUsed]] = {
+      for {
+        dataType ← (artifact \ "dataType").asOpt[String]
+        d ← (artifact \ "data").asOpt[String]
+        data ← (dataType, d) match {
+          case ("file", dataExtractor(filename, contentType, b64content)) ⇒
+            val content = java.util.Base64.getDecoder.decode(b64content)
+            val hashes = Hasher(hashAlg: _*).fromByteArray(content)
+            Some(Right(Attachment(filename, hashes, content.length.toLong, contentType, "")))
+          case ("file", _) ⇒
+            logger.warn(s"Invalid data format for file artifact: $d")
+            None
+          case _ ⇒
+            Some(Left(d))
+        }
+      } 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 \ "0" \ "count").asOpt[Int].getOrElse(0)
+      } 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, 0)
+        case a            ⇒ (a.parentId, 0, 1)
+      }
+      .reduce[(Option[String], Int, Int)] {
+        case ((caze, iocCount1, artifactCount1), (_, iocCount2, artifactCount2)) ⇒ (caze, 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 _ ⇒ Future.failed(InternalError("Case not found"))
+      }
+      .runWith(Sink.seq)
+  }
 }

thehive-backend/app/services/ArtifactSrv.scala

@@ -4,19 +4,17 @@ import javax.inject.{ Inject, Singleton }
 
 import akka.NotUsed
 import akka.stream.scaladsl.Source
-
-import scala.concurrent.{ ExecutionContext, Future }
-import scala.util.{ Failure, Try }
-import play.api.libs.json.JsValue.jsValueToJsLookup
-import org.elastic4play.{ ConflictError, CreateError }
+import models.{ CaseResolutionStatus, CaseStatus, _ }
+import org.elastic4play.ConflictError
 import org.elastic4play.controllers.Fields
-import org.elastic4play.services.{ Agg, AuthContext, CreateSrv, DeleteSrv, FieldsSrv, FindSrv, GetSrv, QueryDSL, QueryDef, UpdateSrv }
-import models.{ Artifact, ArtifactModel, ArtifactStatus, Case, CaseModel }
+import org.elastic4play.services._
 import org.elastic4play.utils.{ RichFuture, RichOr }
-import models.CaseStatus
-import models.CaseResolutionStatus
 import play.api.Logger
 import play.api.libs.json.JsObject
+import play.api.libs.json.JsValue.jsValueToJsLookup
+
+import scala.concurrent.{ ExecutionContext, Future }
+import scala.util.{ Failure, Try }
 
 @Singleton
 class ArtifactSrv @Inject() (
@@ -99,36 +97,54 @@ class ArtifactSrv @Inject() (
     }
   }
 
-  def findSimilar(artifact: Artifact, range: Option[String], sortBy: Seq[String]): (Source[Artifact, NotUsed], Future[Long]) =
+  def findSimilar(artifact: Artifact, range: Option[String], sortBy: Seq[String]): (Source[Artifact, NotUsed], Future[Long]) = {
     find(similarArtifactFilter(artifact), range, sortBy)
+  }
+
+  def findSimilar(dataType: String, data: Either[String, Attachment], filter: Option[QueryDef], range: Option[String], sortBy: Seq[String]): (Source[Artifact, NotUsed], Future[Long]) = {
+    find(similarArtifactFilter(dataType, data, filter.getOrElse(org.elastic4play.services.QueryDSL.any)), range, sortBy)
+  }
 
   private[services] def similarArtifactFilter(artifact: Artifact): QueryDef = {
     import org.elastic4play.services.QueryDSL._
-    val dataType = artifact.dataType()
-    artifact.data() match {
+    val data = (artifact.data(), artifact.attachment()) match {
+      case (Some(_data), None)      ⇒ Left(_data)
+      case (None, Some(attachment)) ⇒ Right(attachment)
+      case _                        ⇒ sys.error("")
+    }
+    val filter = parent("case", not(withId(artifact.parentId.get)))
+    similarArtifactFilter(artifact.dataType(), data, filter)
+  }
+
+  private[services] def similarArtifactFilter(dataType: String, data: Either[String, Attachment], filter: QueryDef): QueryDef = {
+    import org.elastic4play.services.QueryDSL._
+    data match {
       // artifact is an hash
-      case Some(d) if dataType == "hash" ⇒
+      case Left(d) if dataType == "hash" ⇒
         and(
-          parent("case", and(not(withId(artifact.parentId.get)), "status" ~!= CaseStatus.Deleted, "resolutionStatus" ~!= CaseResolutionStatus.Duplicated)),
+          filter,
+          parent("case", and("status" ~!= CaseStatus.Deleted, "resolutionStatus" ~!= CaseResolutionStatus.Duplicated)),
           "status" ~= "Ok",
           or(
             and(
               "data" ~= d,
               "dataType" ~= dataType),
             "attachment.hashes" ~= d))
       // artifact contains data but not an hash
-      case Some(d) ⇒
+      case Left(d) ⇒
         and(
-          parent("case", and(not(withId(artifact.parentId.get)), "status" ~!= CaseStatus.Deleted, "resolutionStatus" ~!= CaseResolutionStatus.Duplicated)),
+          filter,
+          parent("case", and("status" ~!= CaseStatus.Deleted, "resolutionStatus" ~!= CaseResolutionStatus.Duplicated)),
           "status" ~= "Ok",
           "data" ~= d,
           "dataType" ~= dataType)
       // artifact is a file
-      case None ⇒
-        val hashes = artifact.attachment().toSeq.flatMap(_.hashes).map(_.toString)
+      case Right(attachment) ⇒
+        val hashes = attachment.hashes.map(_.toString)
         val hashFilter = hashes.map { h ⇒ "attachment.hashes" ~= h }
         and(
-          parent("case", and(not(withId(artifact.parentId.get)), "status" ~!= CaseStatus.Deleted, "resolutionStatus" ~!= CaseResolutionStatus.Duplicated)),
+          filter,
+          parent("case", and("status" ~!= CaseStatus.Deleted, "resolutionStatus" ~!= CaseResolutionStatus.Duplicated)),
           "status" ~= "Ok",
           or(
             hashFilter :+

thehive-backend/app/services/JsonFormat.scala

@@ -0,0 +1,24 @@
+package services
+
+import play.api.libs.json.{ Json, OWrites }
+
+object JsonFormat {
+  implicit val caseSimilarityWrites: OWrites[CaseSimilarity] = OWrites[CaseSimilarity] {
+    case CaseSimilarity(caze, similarIocCount, iocCount, similarArtifactCount, artifactCount) ⇒
+      Json.obj(
+        "id" → caze.id,
+        "_id" → caze.id,
+        "title" → caze.title(),
+        "tags" → caze.tags(),
+        "status" → caze.status(),
+        "severity" → caze.severity(),
+        "resolutionStatus" → caze.resolutionStatus(),
+        "tlp" → caze.tlp(),
+        "startDate" → caze.startDate(),
+        "endDate" → caze.endDate(),
+        "similarIocCount" → similarIocCount,
+        "iocCount" → iocCount,
+        "similarArtifactCount" → similarArtifactCount,
+        "artifactCount" → artifactCount)
+  }
+}