#1410 Add caseId in alerts

TheHive-Project · Jun 28, 2020 · e0e73b4 · e0e73b4
1 parent 46f67dc
commit e0e73b4
Show file tree

Hide file tree

Showing 2 changed files with 87 additions and 5 deletions.
diff --git a/dto/src/main/scala/org/thp/thehive/dto/v1/Alert.scala b/dto/src/main/scala/org/thp/thehive/dto/v1/Alert.scala
@@ -3,7 +3,7 @@ package org.thp.thehive.dto.v1
 import java.util.Date
 
 import org.thp.scalligraph.controllers.WithParser
-import play.api.libs.json.{Json, OFormat, OWrites}
+import play.api.libs.json.{Json, OFormat, OWrites, Reads}
 
 case class InputAlert(
     `type`: String,
@@ -48,9 +48,89 @@ case class OutputAlert(
     follow: Boolean,
     customFields: Set[OutputCustomFieldValue] = Set.empty,
     caseTemplate: Option[String] = None,
-    observableCount: Long
+    observableCount: Long,
+    caseId: Option[String]
 )
 
 object OutputAlert {
-  implicit val format: OFormat[OutputAlert] = Json.format[OutputAlert]
+  implicit val reads: Reads[OutputAlert] = Reads[OutputAlert] { json =>
+    for {
+      _id             <- (json \ "_id").validate[String]
+      _type           <- (json \ "_type").validate[String]
+      _createdBy      <- (json \ "_createdBy").validate[String]
+      _updatedBy      <- (json \ "_updatedBy").validateOpt[String]
+      _createdAt      <- (json \ "_createdAt").validate[Date]
+      _updatedAt      <- (json \ "_updatedAt").validateOpt[Date]
+      tpe             <- (json \ "type").validate[String]
+      source          <- (json \ "source").validate[String]
+      sourceRef       <- (json \ "sourceRef").validate[String]
+      externalLink    <- (json \ "externalLink").validateOpt[String]
+      title           <- (json \ "title").validate[String]
+      description     <- (json \ "description").validate[String]
+      severity        <- (json \ "severity").validate[Int]
+      date            <- (json \ "date").validate[Date]
+      tags            <- (json \ "tags").validate[Set[String]]
+      tlp             <- (json \ "tlp").validate[Int]
+      pap             <- (json \ "pap").validate[Int]
+      read            <- (json \ "read").validate[Boolean]
+      follow          <- (json \ "follow").validate[Boolean]
+      customFields    <- (json \ "customFields").validate[Set[OutputCustomFieldValue]]
+      caseTemplate    <- (json \ "caseTemplate").validateOpt[String]
+      observableCount <- (json \ "observableCount").validate[Long]
+      caseId          <- (json \ "caseId").validateOpt[String]
+    } yield OutputAlert(
+      _id,
+      _type,
+      _createdBy,
+      _updatedBy,
+      _createdAt,
+      _updatedAt,
+      tpe,
+      source,
+      sourceRef,
+      externalLink,
+      title,
+      description,
+      severity,
+      date,
+      tags,
+      tlp,
+      pap,
+      read,
+      follow,
+      customFields,
+      caseTemplate,
+      observableCount,
+      caseId
+    )
+  }
+  implicit val writes: OWrites[OutputAlert] = OWrites[OutputAlert] { outputAlert =>
+    Json.obj(
+      "_id"             -> outputAlert._id,
+      "_type"           -> outputAlert._type,
+      "_createdBy"      -> outputAlert._createdBy,
+      "_updatedBy"      -> outputAlert._updatedBy,
+      "_createdAt"      -> outputAlert._createdAt,
+      "_updatedAt"      -> outputAlert._updatedAt,
+      "type"            -> outputAlert.`type`,
+      "source"          -> outputAlert.source,
+      "sourceRef"       -> outputAlert.sourceRef,
+      "externalLink"    -> outputAlert.externalLink,
+      "title"           -> outputAlert.title,
+      "description"     -> outputAlert.description,
+      "severity"        -> outputAlert.severity,
+      "date"            -> outputAlert.date,
+      "tags"            -> outputAlert.tags,
+      "tlp"             -> outputAlert.tlp,
+      "pap"             -> outputAlert.pap,
+      "read"            -> outputAlert.read,
+      "follow"          -> outputAlert.follow,
+      "customFields"    -> outputAlert.customFields,
+      "caseTemplate"    -> outputAlert.caseTemplate,
+      "observableCount" -> outputAlert.observableCount,
+      "caseId"          -> outputAlert.caseId
+    )
+  }
+
+  implicit val format: OFormat[OutputAlert] = OFormat(reads, writes)
 }
diff --git a/thehive/test/org/thp/thehive/controllers/v1/AlertCtrlTest.scala b/thehive/test/org/thp/thehive/controllers/v1/AlertCtrlTest.scala
@@ -77,7 +77,8 @@ class AlertCtrlTest extends PlaySpecification with TestAppBuilder {
         follow = true,
         customFields = Set.empty,
         caseTemplate = None,
-        observableCount = 0L
+        observableCount = 0L,
+        caseId = None
       )
 
       createdAlert must_=== expected
@@ -123,7 +124,8 @@ class AlertCtrlTest extends PlaySpecification with TestAppBuilder {
         follow = true,
         customFields = Set.empty,
         caseTemplate = Some("spam"),
-        observableCount = 0L
+        observableCount = 0L,
+        caseId = None
       )
 
       createdAlert must_=== expected