#1636 Fix customField filter in alert

TheHive-Project · Nov 10, 2020 · c090915 · c090915
1 parent 210f383
commit c090915
Show file tree

Hide file tree

Showing 5 changed files with 64 additions and 21 deletions.
diff --git a/ScalliGraph b/ScalliGraph
diff --git a/thehive/app/org/thp/thehive/controllers/v0/AlertCtrl.scala b/thehive/app/org/thp/thehive/controllers/v0/AlertCtrl.scala
@@ -423,7 +423,7 @@ class PublicAlert @Inject() (
       .property("user", UMapping.string)(_.field.updatable)
       .property("customFields", UMapping.jsonNative)(_.subSelect {
         case (FPathElem(_, FPathElem(name, _)), alertSteps) =>
-          alertSteps.customFields(name).jsonValue
+          alertSteps.customFields(EntityIdOrName(name)).jsonValue
         case (_, alertSteps) => alertSteps.customFields.nameJsonValue.fold.domainMap(JsObject(_))
       }.custom {
         case (FPathElem(_, FPathElem(name, _)), value, vertex, _, graph, authContext) =>

diff --git a/thehive/app/org/thp/thehive/controllers/v1/Properties.scala b/thehive/app/org/thp/thehive/controllers/v1/Properties.scala
@@ -10,7 +10,6 @@ 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._
@@ -93,16 +92,56 @@ class Properties @Inject() (
       .property("summary", UMapping.string.optional)(_.field.updatable)
       .property("user", UMapping.string)(_.field.updatable)
       .property("customFields", UMapping.jsonNative)(_.subSelect {
-        case (FPathElem(_, FPathElem(name, _)), alertSteps) => alertSteps.customFields(name).jsonValue
-        case (_, alertSteps)                                => alertSteps.customFields.nameJsonValue.fold.domainMap(JsObject(_))
-      }.custom {
-        case (FPathElem(_, FPathElem(name, _)), value, vertex, _, graph, authContext) =>
-          for {
-            c <- alertSrv.getOrFail(vertex)(graph)
-            _ <- alertSrv.setOrCreateCustomField(c, InputCustomFieldValue(name, Some(value), None))(graph, authContext)
-          } yield Json.obj(s"customField.$name" -> value)
-        case _ => Failure(BadRequestError("Invalid custom fields format"))
-      })
+        case (FPathElem(_, FPathElem(idOrName, _)), alerts) =>
+          alerts
+            .customFields(EntityIdOrName(idOrName))
+            .jsonValue
+        case (_, caseSteps) => caseSteps.customFields.nameJsonValue.fold.domainMap(JsObject(_))
+      }
+        .filter {
+          case (FPathElem(_, FPathElem(idOrName, _)), caseTraversal) =>
+            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)
+              }
+              .getOrElse(caseTraversal.constant2(null))
+          case (_, caseTraversal) => caseTraversal.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(new Converter[Any, JsValue] { def apply(x: JsValue): Any = x })
+          case _ => (x: JsValue) => x
+        }
+        .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)
+            } 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)
+            } yield Json.obj("customFields" -> values)
+          case _ => Failure(BadRequestError("Invalid custom fields format"))
+        })
       .build
 
   lazy val audit: PublicProperties =

diff --git a/thehive/app/org/thp/thehive/services/AlertSrv.scala b/thehive/app/org/thp/thehive/services/AlertSrv.scala
@@ -183,15 +183,15 @@ class AlertSrv @Inject() (
       graph: Graph,
       authContext: AuthContext
   ): Try[Unit] = {
-    val cfv = get(alert).customFields(cf.name)
+    val cfv = get(alert).customFields(EntityIdOrName(cf.name))
     if (cfv.clone().exists)
       cfv.setValue(cf.value)
     else
       createCustomField(alert, cf).map(_ => ())
   }
 
-  def getCustomField(alert: Alert with Entity, customFieldName: String)(implicit graph: Graph): Option[RichCustomField] =
-    get(alert).customFields(customFieldName).richCustomField.headOption
+//  def getCustomField(alert: Alert with Entity, customFieldName: String)(implicit graph: Graph): Option[RichCustomField] =
+//    get(alert).customFields(customFieldName).richCustomField.headOption
 
   def updateCustomField(
       alert: Alert with Entity,
@@ -203,7 +203,7 @@ class AlertSrv @Inject() (
       .richCustomField
       .toIterator
       .filterNot(rcf => customFieldNames.contains(rcf.name))
-      .foreach(rcf => get(alert).customFields(rcf.name).remove())
+      .foreach(rcf => get(alert).customFields(rcf.customField._id).remove())
     customFieldValues
       .toTry { case (cf, v) => setOrCreateCustomField(alert, InputCustomFieldValue(cf.name, Some(v), None)) }
       .map(_ => ())
@@ -494,8 +494,12 @@ object AlertOps {
         result.limit(0)
     }
 
-    def customFields(name: String): Traversal.E[AlertCustomField] =
-      traversal.outE[AlertCustomField].filter(_.inV.v[CustomField].has(_.name, name))
+    def customFields(idOrName: EntityIdOrName): Traversal.E[AlertCustomField] =
+      idOrName
+        .fold(
+          id => traversal.outE[AlertCustomField].filter(_.inV.getByIds(id)),
+          name => traversal.outE[AlertCustomField].filter(_.inV.v[CustomField].has(_.name, name))
+        )
 
     def customFields: Traversal.E[AlertCustomField] = traversal.outE[AlertCustomField]
 

diff --git a/thehive/test/org/thp/thehive/services/AlertSrvTest.scala b/thehive/test/org/thp/thehive/services/AlertSrvTest.scala
@@ -2,7 +2,7 @@ package org.thp.thehive.services
 
 import java.util.Date
 
-import org.thp.scalligraph.EntityName
+import org.thp.scalligraph.{EntityIdOrName, EntityName}
 import org.thp.scalligraph.auth.AuthContext
 import org.thp.scalligraph.models._
 import org.thp.scalligraph.traversal.TraversalOps._
@@ -149,7 +149,7 @@ class AlertSrvTest extends PlaySpecification with TestAppBuilder {
       } must beSuccessfulTry
 
       app[Database].roTransaction { implicit graph =>
-        app[AlertSrv].get(EntityName("testType;testSource;ref1")).customFields("string1").nameJsonValue.headOption
+        app[AlertSrv].get(EntityName("testType;testSource;ref1")).customFields(EntityIdOrName("string1")).nameJsonValue.headOption
       } must beSome("string1" -> JsString("sad"))
     }