From f36c1474fc193583d0116e69475d2100e067b74f Mon Sep 17 00:00:00 2001
From: To-om <toom@thehive-project.org>
Date: Mon, 23 Oct 2017 13:30:54 +0200
Subject: [PATCH] #137 Add "seen" attribute in alert artifact

---
 .../app/controllers/AlertCtrl.scala            |  6 +++++-
 thehive-backend/app/services/AlertSrv.scala    | 18 ++++++++++++++++++
 2 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/thehive-backend/app/controllers/AlertCtrl.scala b/thehive-backend/app/controllers/AlertCtrl.scala
index 322f90f3dd..c0cea8cf3e 100644
--- 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
index e5baa8743f..4759ddaafa 100644
--- 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._