#1779 Fix custom fields property definition

TheHive-Project · Feb 5, 2021 · e1fd7ae · e1fd7ae
1 parent 42d554c
commit e1fd7ae
Show file tree

Hide file tree

Showing 2 changed files with 66 additions and 36 deletions.
diff --git a/thehive/app/org/thp/thehive/controllers/v0/AlertCtrl.scala b/thehive/app/org/thp/thehive/controllers/v0/AlertCtrl.scala
@@ -21,10 +21,10 @@ import org.thp.thehive.services.OrganisationOps._
 import org.thp.thehive.services.TagOps._
 import org.thp.thehive.services.UserOps._
 import org.thp.thehive.services._
-import play.api.libs.json.{JsArray, JsObject, Json}
+import play.api.libs.json.{JsArray, JsObject, JsValue, Json}
 import play.api.mvc.{Action, AnyContent, Results}
 
-import java.util.{Base64, List => JList, Map => JMap}
+import java.util.{Base64, Date, List => JList, Map => JMap}
 import javax.inject.{Inject, Named, Singleton}
 import scala.util.{Failure, Success, Try}
 
@@ -433,21 +433,52 @@ class PublicAlert @Inject() (
         case (FPathElem(_, FPathElem(name, _)), alertSteps) =>
           alertSteps.customFields(EntityIdOrName(name)).jsonValue
         case (_, alertSteps) => alertSteps.customFields.nameJsonValue.fold.domainMap(JsObject(_))
-      }.custom {
-        case (FPathElem(_, FPathElem(name, _)), value, vertex, _, graph, authContext) =>
-          for {
-            c <- alertSrv.getByIds(EntityId(vertex.id))(graph).getOrFail("Alert")
-            _ <- alertSrv.setOrCreateCustomField(c, InputCustomFieldValue(name, Some(value), None))(graph, authContext)
-          } yield Json.obj(s"customField.$name" -> value)
-        case (FPathElem(_, FPathEmpty), values: JsObject, vertex, _, graph, authContext) =>
-          for {
-            c   <- alertSrv.get(vertex)(graph).getOrFail("Alert")
-            cfv <- values.fields.toTry { case (n, v) => customFieldSrv.getOrFail(EntityIdOrName(n))(graph).map(_ -> v) }
-            _   <- alertSrv.updateCustomField(c, cfv)(graph, authContext)
-          } yield Json.obj("customFields" -> values)
+      }
+        .filter {
+          case (FPathElem(_, FPathElem(idOrName, _)), alerts) =>
+            db
+              .roTransaction(implicit graph => customFieldSrv.get(EntityIdOrName(idOrName)).value(_.`type`).getOrFail("CustomField"))
+              .map {
+                case CustomFieldType.boolean => alerts.customFields(EntityIdOrName(idOrName)).value(_.booleanValue)
+                case CustomFieldType.date    => alerts.customFields(EntityIdOrName(idOrName)).value(_.dateValue)
+                case CustomFieldType.float   => alerts.customFields(EntityIdOrName(idOrName)).value(_.floatValue)
+                case CustomFieldType.integer => alerts.customFields(EntityIdOrName(idOrName)).value(_.integerValue)
+                case CustomFieldType.string  => alerts.customFields(EntityIdOrName(idOrName)).value(_.stringValue)
+              }
+              .getOrElse(alerts.constant2(null))
+          case (_, alerts) => alerts.constant2(null)
+        }
+        .converter {
+          case FPathElem(_, FPathElem(idOrName, _)) =>
+            db
+              .roTransaction { implicit graph =>
+                customFieldSrv.get(EntityIdOrName(idOrName)).value(_.`type`).getOrFail("CustomField")
+              }
+              .map {
+                case CustomFieldType.boolean => new Converter[Any, JsValue] { def apply(x: JsValue): Any = x.as[Boolean] }
+                case CustomFieldType.date    => new Converter[Any, JsValue] { def apply(x: JsValue): Any = x.as[Date] }
+                case CustomFieldType.float   => new Converter[Any, JsValue] { def apply(x: JsValue): Any = x.as[Double] }
+                case CustomFieldType.integer => new Converter[Any, JsValue] { def apply(x: JsValue): Any = x.as[Long] }
+                case CustomFieldType.string  => new Converter[Any, JsValue] { def apply(x: JsValue): Any = x.as[String] }
+              }
+              .getOrElse((x: JsValue) => x)
+          case _ => (x: JsValue) => x
+        }
+        .custom {
+          case (FPathElem(_, FPathElem(name, _)), value, vertex, _, graph, authContext) =>
+            for {
+              c <- alertSrv.getByIds(EntityId(vertex.id))(graph).getOrFail("Alert")
+              _ <- alertSrv.setOrCreateCustomField(c, InputCustomFieldValue(name, Some(value), None))(graph, authContext)
+            } yield Json.obj(s"customField.$name" -> value)
+          case (FPathElem(_, FPathEmpty), values: JsObject, vertex, _, graph, authContext) =>
+            for {
+              c   <- alertSrv.get(vertex)(graph).getOrFail("Alert")
+              cfv <- values.fields.toTry { case (n, v) => customFieldSrv.getOrFail(EntityIdOrName(n))(graph).map(_ -> v) }
+              _   <- alertSrv.updateCustomField(c, cfv)(graph, authContext)
+            } yield Json.obj("customFields" -> values)
 
-        case _ => Failure(BadRequestError("Invalid custom fields format"))
-      })
+          case _ => Failure(BadRequestError("Invalid custom fields format"))
+        })
       .property("case", db.idMapping)(_.select(_.`case`._id).readonly)
       .property("imported", UMapping.boolean)(_.select(_.imported).readonly)
       .property("importDate", UMapping.date.optional)(_.select(_.importDate).readonly)

diff --git a/thehive/app/org/thp/thehive/controllers/v1/Properties.scala b/thehive/app/org/thp/thehive/controllers/v1/Properties.scala
@@ -1,15 +1,12 @@
 package org.thp.thehive.controllers.v1
 
-import java.lang.{Long => JLong}
-import java.util.Date
-
-import javax.inject.{Inject, Named, Singleton}
 import org.thp.scalligraph.controllers.{FPathElem, FPathEmpty}
 import org.thp.scalligraph.models.{Database, UMapping}
 import org.thp.scalligraph.query.{PublicProperties, PublicPropertyListBuilder}
 import org.thp.scalligraph.traversal.Converter
 import org.thp.scalligraph.traversal.TraversalOps._
 import org.thp.scalligraph.{BadRequestError, EntityIdOrName, RichSeq}
+import org.thp.thehive.dto.v1.InputCustomFieldValue
 import org.thp.thehive.models._
 import org.thp.thehive.services.AlertOps._
 import org.thp.thehive.services.AuditOps._
@@ -25,6 +22,8 @@ import org.thp.thehive.services.UserOps._
 import org.thp.thehive.services._
 import play.api.libs.json.{JsObject, JsValue, Json}
 
+import java.util.Date
+import javax.inject.{Inject, Named, Singleton}
 import scala.util.Failure
 
 @Singleton
@@ -60,8 +59,8 @@ class Properties @Inject() (
       .property("lastSyncDate", UMapping.date.optional)(_.field.updatable)
       .property("tags", UMapping.string.set)(
         _.select(_.tags.displayName)
-          .filter((_, cases) =>
-            cases
+          .filter((_, alerts) =>
+            alerts
               .tags
               .graphMap[String, String, Converter.Identity[String]](
                 { v =>
@@ -96,21 +95,21 @@ class Properties @Inject() (
           alerts
             .customFields(EntityIdOrName(idOrName))
             .jsonValue
-        case (_, caseSteps) => caseSteps.customFields.nameJsonValue.fold.domainMap(JsObject(_))
+        case (_, alerts) => alerts.customFields.nameJsonValue.fold.domainMap(JsObject(_))
       }
         .filter {
-          case (FPathElem(_, FPathElem(idOrName, _)), caseTraversal) =>
+          case (FPathElem(_, FPathElem(idOrName, _)), alerts) =>
             db
               .roTransaction(implicit graph => customFieldSrv.get(EntityIdOrName(idOrName)).value(_.`type`).getOrFail("CustomField"))
               .map {
-                case CustomFieldType.boolean => caseTraversal.customFields(EntityIdOrName(idOrName)).value(_.booleanValue)
-                case CustomFieldType.date    => caseTraversal.customFields(EntityIdOrName(idOrName)).value(_.dateValue)
-                case CustomFieldType.float   => caseTraversal.customFields(EntityIdOrName(idOrName)).value(_.floatValue)
-                case CustomFieldType.integer => caseTraversal.customFields(EntityIdOrName(idOrName)).value(_.integerValue)
-                case CustomFieldType.string  => caseTraversal.customFields(EntityIdOrName(idOrName)).value(_.stringValue)
+                case CustomFieldType.boolean => alerts.customFields(EntityIdOrName(idOrName)).value(_.booleanValue)
+                case CustomFieldType.date    => alerts.customFields(EntityIdOrName(idOrName)).value(_.dateValue)
+                case CustomFieldType.float   => alerts.customFields(EntityIdOrName(idOrName)).value(_.floatValue)
+                case CustomFieldType.integer => alerts.customFields(EntityIdOrName(idOrName)).value(_.integerValue)
+                case CustomFieldType.string  => alerts.customFields(EntityIdOrName(idOrName)).value(_.stringValue)
               }
-              .getOrElse(caseTraversal.constant2(null))
-          case (_, caseTraversal) => caseTraversal.constant2(null)
+              .getOrElse(alerts.constant2(null))
+          case (_, alerts) => alerts.constant2(null)
         }
         .converter {
           case FPathElem(_, FPathElem(idOrName, _)) =>
@@ -131,14 +130,14 @@ class Properties @Inject() (
         .custom {
           case (FPathElem(_, FPathElem(idOrName, _)), value, vertex, _, graph, authContext) =>
             for {
-              c <- caseSrv.get(vertex)(graph).getOrFail("Case")
-              _ <- caseSrv.setOrCreateCustomField(c, EntityIdOrName(idOrName), Some(value), None)(graph, authContext)
+              a <- alertSrv.get(vertex)(graph).getOrFail("Alert")
+              _ <- alertSrv.setOrCreateCustomField(a, InputCustomFieldValue(idOrName, Some(value), None))(graph, authContext)
             } yield Json.obj(s"customField.$idOrName" -> value)
           case (FPathElem(_, FPathEmpty), values: JsObject, vertex, _, graph, authContext) =>
             for {
-              c   <- caseSrv.get(vertex)(graph).getOrFail("Case")
-              cfv <- values.fields.toTry { case (n, v) => customFieldSrv.getOrFail(EntityIdOrName(n))(graph).map(cf => (cf, v, None)) }
-              _   <- caseSrv.updateCustomField(c, cfv)(graph, authContext)
+              c   <- alertSrv.get(vertex)(graph).getOrFail("Alert")
+              cfv <- values.fields.toTry { case (n, v) => customFieldSrv.getOrFail(EntityIdOrName(n))(graph).map(_ -> v) }
+              _   <- alertSrv.updateCustomField(c, cfv)(graph, authContext)
             } yield Json.obj("customFields" -> values)
           case _ => Failure(BadRequestError("Invalid custom fields format"))
         })