Skip to content

Commit

Permalink
#237 Save alert attachment in datastore
Browse files Browse the repository at this point in the history
  • Loading branch information
@To-om
To-om committed Jun 22, 2017
1 parent 959b5b2 commit a8d3cc5
Show file tree
Hide file tree
Showing 3 changed files with 62 additions and 26 deletions.
32 changes: 25 additions & 7 deletions thehive-backend/app/models/Alert.scala
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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 {
Expand All @@ -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)))
Expand Down
23 changes: 20 additions & 3 deletions thehive-backend/app/services/AlertSrv.scala
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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]
}

Expand Down Expand Up @@ -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)
Expand Down
33 changes: 17 additions & 16 deletions thehive-misp/app/connectors/misp/MispSrv.scala
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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))
Expand All @@ -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 {
Expand All @@ -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)
Expand Down

0 comments on commit a8d3cc5

Please sign in to comment.