#137 Add "seen" attribute in alert artifact

TheHive-Project · Oct 23, 2017 · f36c147 · f36c147
1 parent 2f82d50
commit f36c147
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 1 deletion.
diff --git a/thehive-backend/app/controllers/AlertCtrl.scala b/thehive-backend/app/controllers/AlertCtrl.scala
@@ -75,8 +75,12 @@ class AlertCtrl @Inject() (
         alertSrv.similarCases(alert)
           .map(sc ⇒ Json.obj("similarCases" → Json.toJson(sc)))
       else Future.successful(JsObject(Nil))
+      similarArtifacts ← if (withSimilarity)
+        alertSrv.alertArtifactsWithSeen(alert)
+          .map(aws ⇒ Json.obj("artifacts" → aws))
+      else Future.successful(JsObject(Nil))
     } yield {
-      renderer.toOutput(OK, alertsWithStats ++ similarCases)
+      renderer.toOutput(OK, alertsWithStats ++ similarCases ++ similarArtifacts)
     }
   }
 

diff --git a/thehive-backend/app/services/AlertSrv.scala b/thehive-backend/app/services/AlertSrv.scala
@@ -320,6 +320,24 @@ class AlertSrv(
       .runWith(Sink.seq)
   }
 
+  def getArtifactSeen(artifact: JsObject): Future[Long] = {
+    val maybeArtifactSeen = for {
+      dataType ← (artifact \ "dataType").asOpt[String]
+      data ← dataType match {
+        case "file" ⇒ (artifact \ "attachment").asOpt[Attachment].map(Right.apply)
+        case _      ⇒ (artifact \ "data").asOpt[String].map(Left.apply)
+      }
+      numberOfSimilarArtifacts = artifactSrv.findSimilar(dataType, data, None, None, Nil)._2
+    } yield numberOfSimilarArtifacts
+    maybeArtifactSeen.getOrElse(Future.successful(0L))
+  }
+
+  def alertArtifactsWithSeen(alert: Alert): Future[Seq[JsObject]] = {
+    Future.traverse(alert.artifacts()) { artifact ⇒
+      getArtifactSeen(artifact).map(seen ⇒ artifact + ("seen" → JsNumber(seen)))
+    }
+  }
+
   def fixStatus()(implicit authContext: AuthContext): Future[Unit] = {
     import org.elastic4play.services.QueryDSL._