#237 Save alert attachment in datastore

Author: To-om

thehive-backend/app/models/Alert.scala

@@ -5,9 +5,10 @@ import javax.inject.{ Inject, Singleton }
 
 import models.JsonFormat.alertStatusFormat
 import org.elastic4play.controllers.JsonInputValue
-import org.elastic4play.{ AttributeCheckingError, InvalidFormatAttributeError }
-import org.elastic4play.models.{ Attribute, AttributeDef, BaseEntity, EntityDef, HiveEnumeration, ModelDef, AttributeFormat ⇒ F, AttributeOption ⇒ O }
+import org.elastic4play.models.{ Attribute, AttributeDef, BaseEntity, EntityDef, HiveEnumeration, ModelDef, MultiAttributeFormat, OptionalAttributeFormat, AttributeFormat ⇒ F, AttributeOption ⇒ O }
+import org.elastic4play.services.DBLists
 import org.elastic4play.utils.Hasher
+import org.elastic4play.{ AttributeCheckingError, InvalidFormatAttributeError }
 import play.api.Logger
 import play.api.libs.json._
 import services.AuditedModel
@@ -43,7 +44,7 @@ trait AlertAttributes {
 }
 
 @Singleton
-class AlertModel @Inject() (artifactModel: ArtifactModel)
+class AlertModel @Inject() (dblists: DBLists)
     extends ModelDef[AlertModel, Alert]("alert")
     with AlertAttributes
     with AuditedModel {
@@ -52,17 +53,34 @@ class AlertModel @Inject() (artifactModel: ArtifactModel)
   override val defaultSortBy: Seq[String] = Seq("-date")
   override val removeAttribute: JsObject = Json.obj("status" → AlertStatus.Ignored)
 
-  override def artifactAttributes: Seq[Attribute[_]] = artifactModel
-    .attributes
-    .filter(_.isForm)
+  override def artifactAttributes: Seq[Attribute[_]] = {
+    val remoteAttachmentAttributes = Seq(
+      Attribute("alert", "reference", F.stringFmt, Nil, None, ""),
+      Attribute("alert", "filename", OptionalAttributeFormat(F.stringFmt), Nil, None, ""),
+      Attribute("alert", "contentType", OptionalAttributeFormat(F.stringFmt), Nil, None, ""),
+      Attribute("alert", "size", OptionalAttributeFormat(F.numberFmt), Nil, None, ""),
+      Attribute("alert", "hash", MultiAttributeFormat(F.stringFmt), Nil, None, ""),
+      Attribute("alert", "type", OptionalAttributeFormat(F.stringFmt), Nil, None, ""))
+
+    Seq(
+      Attribute("alert", "data", OptionalAttributeFormat(F.stringFmt), Nil, None, ""),
+      Attribute("alert", "dataType", F.stringFmt, Nil, None, ""),
+      Attribute("alert", "message", OptionalAttributeFormat(F.stringFmt), Nil, None, ""),
+      Attribute("alert", "startDate", OptionalAttributeFormat(F.dateFmt), Nil, None, ""),
+      Attribute("alert", "attachment", OptionalAttributeFormat(F.attachmentFmt), Nil, None, ""),
+      Attribute("alert", "remoteAttachment", OptionalAttributeFormat(F.objectFmt(remoteAttachmentAttributes)), Nil, None, ""),
+      Attribute("alert", "tlp", OptionalAttributeFormat(F.numberFmt), Nil, None, ""),
+      Attribute("alert", "tags", MultiAttributeFormat(F.stringFmt), Nil, None, ""),
+      Attribute("alert", "ioc", OptionalAttributeFormat(F.stringFmt), Nil, None, ""))
+  }
 
   override def creationHook(parent: Option[BaseEntity], attrs: JsObject): Future[JsObject] = {
     // check if data attribute is present on all artifacts
     val missingDataErrors = (attrs \ "artifacts")
       .asOpt[Seq[JsValue]]
       .getOrElse(Nil)
       .filter { a ⇒
-        (a \ "data").toOption.isEmpty ||
+        ((a \ "data").toOption.isEmpty && (a \ "attachment").toOption.isEmpty && (a \ "remoteAttachment").toOption.isEmpty) ||
           ((a \ "tags").toOption.isEmpty && (a \ "message").toOption.isEmpty)
       }
       .map(v ⇒ InvalidFormatAttributeError("artifacts", "artifact", JsonInputValue(v)))

thehive-backend/app/services/AlertSrv.scala

@@ -17,10 +17,12 @@ import play.api.{ Configuration, Logger }
 
 import scala.collection.immutable
 import scala.concurrent.{ ExecutionContext, Future }
-import scala.util.{ Failure, Success, Try }
+import scala.util.{ Failure, Try }
+import org.elastic4play.services.JsonFormat.attachmentFormat
 
 trait AlertTransformer {
   def createCase(alert: Alert, customCaseTemplate: Option[String])(implicit authContext: AuthContext): Future[Case]
+
   def mergeWithCase(alert: Alert, caze: Case)(implicit authContext: AuthContext): Future[Case]
 }
 
@@ -76,8 +78,23 @@ class AlertSrv(
 
   private[AlertSrv] lazy val logger = Logger(getClass)
 
-  def create(fields: Fields)(implicit authContext: AuthContext): Future[Alert] =
-    createSrv[AlertModel, Alert](alertModel, fields)
+  def create(fields: Fields)(implicit authContext: AuthContext): Future[Alert] = {
+
+    val artifactsFields =
+      Future.traverse(fields.getValues("artifacts")) {
+        case a: JsObject if (a \ "dataType").asOpt[String].contains("file") ⇒
+          (a \ "data").asOpt[String] match {
+            case Some(dataExtractor(filename, contentType, data)) ⇒
+              attachmentSrv.save(filename, contentType, java.util.Base64.getDecoder.decode(data))
+                .map(attachment ⇒ a - "data" + ("attachment" → Json.toJson(attachment)))
+            case _ ⇒ Future.successful(a)
+          }
+        case a ⇒ Future.successful(a)
+      }
+    artifactsFields.flatMap { af ⇒
+      createSrv[AlertModel, Alert](alertModel, fields.set("artifacts", JsArray(af)))
+    }
+  }
 
   def bulkCreate(fieldSet: Seq[Fields])(implicit authContext: AuthContext): Future[Seq[Try[Alert]]] =
     createSrv[AlertModel, Alert](alertModel, fieldSet)

thehive-misp/app/connectors/misp/MispSrv.scala

@@ -300,10 +300,10 @@ class MispSrv @Inject() (
                   "dataType" → "file",
                   "message" → a.comment,
                   "tags" → (artifactTags.value ++ a.tags.map(JsString)),
-                  "data" → Json.obj(
+                  "remoteAttachment" → Json.obj(
                     "filename" → a.value,
-                    "attributeId" → a.id,
-                    "attributeType" → a.tpe).toString,
+                    "reference" → a.id,
+                    "type" → a.tpe),
                   "startDate" → a.date))
             case a ⇒ convertAttribute(a).map { j ⇒
               val tags = artifactTags ++ (j \ "tags").asOpt[JsArray].getOrElse(JsArray(Nil))
@@ -319,20 +319,20 @@ class MispSrv @Inject() (
     attr: JsObject)(implicit authContext: AuthContext): Option[Future[Fields]] = {
     (for {
       dataType ← (attr \ "dataType").validate[String]
-      data ← (attr \ "data").validate[String]
+      data ← (attr \ "data").validateOpt[String]
       message ← (attr \ "message").validate[String]
       startDate ← (attr \ "startDate").validate[Date]
-      attachment = dataType match {
-        case "file" ⇒
-          val json = Json.parse(data)
-          for {
-            attributeId ← (json \ "attributeId").asOpt[String]
-            attributeType ← (json \ "attributeType").asOpt[String]
-            fiv = downloadAttachment(mispConnection, attributeId)
-          } yield if (attributeType == "malware-sample") fiv.map(extractMalwareAttachment)
-          else fiv
-        case _ ⇒ None
-      }
+      attachmentReference ← (attr \ "remoteAttachment" \ "reference").validateOpt[String]
+      attachmentType ← (attr \ "remoteAttachment" \ "type").validateOpt[String]
+      attachment = attachmentReference
+        .flatMap {
+          case ref if dataType == "file" ⇒ Some(downloadAttachment(mispConnection, ref))
+          case _                         ⇒ None
+        }
+        .map {
+          case f if attachmentType.contains("malware-sample") ⇒ f.map(extractMalwareAttachment)
+          case f                                              ⇒ f
+        }
       tags = (attr \ "tags").asOpt[Seq[String]].getOrElse(Nil)
       tlp = tags.map(_.toLowerCase)
         .collectFirst {
@@ -351,7 +351,8 @@ class MispSrv @Inject() (
             .filterNot(_.toLowerCase.startsWith("tlp:"))
             .map(JsString)))
         .set("tlp", tlp)
-    } yield attachment.fold(Future.successful(fields.set("data", data)))(_.map { fiv ⇒
+      if attachment.isDefined != data.isDefined
+    } yield attachment.fold(Future.successful(fields.set("data", data.get)))(_.map { fiv ⇒
       fields.set("attachment", fiv)
     })) match {
       case JsSuccess(r, _) ⇒ Some(r)