#1921 Fix customField selection query

TheHive-Project · Apr 1, 2021 · 14eb881 · 14eb881
1 parent 6e190d3
commit 14eb881
Show file tree

Hide file tree

Showing 9 changed files with 174 additions and 51 deletions.
diff --git a/ScalliGraph b/ScalliGraph
diff --git a/thehive/app/org/thp/thehive/controllers/v0/CaseCtrl.scala b/thehive/app/org/thp/thehive/controllers/v0/CaseCtrl.scala
@@ -275,9 +275,7 @@ class PublicCase @Inject() (
       })
       .property("customFields", UMapping.jsonNative)(_.subSelect {
         case (FPathElem(_, FPathElem(idOrName, _)), caseSteps) =>
-          caseSteps
-            .customFields(EntityIdOrName(idOrName))
-            .jsonValue
+          caseSteps.customFieldJsonValue(customFieldSrv, EntityIdOrName(idOrName))
         case (_, caseSteps) => caseSteps.customFields.nameJsonValue.fold.domainMap(JsObject(_))
       }
         .filter[JsValue] {

diff --git a/thehive/app/org/thp/thehive/controllers/v0/CaseTemplateCtrl.scala b/thehive/app/org/thp/thehive/controllers/v0/CaseTemplateCtrl.scala
@@ -16,7 +16,7 @@ import org.thp.thehive.services.TaskOps._
 import org.thp.thehive.services.UserOps._
 import org.thp.thehive.services._
 import play.api.Logger
-import play.api.libs.json.{JsObject, Json}
+import play.api.libs.json.{JsObject, JsValue, Json}
 import play.api.mvc.{Action, AnyContent, Results}
 
 import javax.inject.{Inject, Named, Singleton}
@@ -134,22 +134,32 @@ class PublicCaseTemplate @Inject() (
     .property("summary", UMapping.string.optional)(_.field.updatable)
     .property("user", UMapping.string)(_.field.updatable)
     .property("customFields", UMapping.jsonNative)(_.subSelect {
-      case (FPathElem(_, FPathElem(name, _)), caseTemplateSteps) => caseTemplateSteps.customFields(name).jsonValue
+      case (FPathElem(_, FPathElem(name, _)), caseTemplateSteps) => caseTemplateSteps.customFieldJsonValue(customFieldSrv, EntityIdOrName(name))
       case (_, caseTemplateSteps)                                => caseTemplateSteps.customFields.nameJsonValue.fold.domainMap(JsObject(_))
-    }.custom {
-      case (FPathElem(_, FPathElem(name, _)), value, vertex, graph, authContext) =>
-        for {
-          c <- caseTemplateSrv.get(vertex)(graph).getOrFail("CaseTemplate")
-          _ <- caseTemplateSrv.setOrCreateCustomField(c, name, Some(value), None)(graph, authContext)
-        } yield Json.obj(s"customFields.$name" -> value)
-      case (FPathElem(_, FPathEmpty), values: JsObject, vertex, graph, authContext) =>
-        for {
-          c   <- caseTemplateSrv.get(vertex)(graph).getOrFail("CaseTemplate")
-          cfv <- values.fields.toTry { case (n, v) => customFieldSrv.getOrFail(EntityIdOrName(n))(graph).map(_ -> v) }
-          _   <- caseTemplateSrv.updateCustomField(c, cfv)(graph, authContext)
-        } yield Json.obj("customFields" -> values)
-      case _ => Failure(BadRequestError("Invalid custom fields format"))
-    })
+    }
+      .filter[JsValue] {
+        case (FPathElem(_, FPathElem(name, _)), caseTemplateTraversal, _, predicate) =>
+          predicate match {
+            case Right(predicate) => caseTemplateTraversal.customFieldFilter(customFieldSrv, EntityIdOrName(name), predicate)
+            case Left(true)       => caseTemplateTraversal.hasCustomField(customFieldSrv, EntityIdOrName(name))
+            case Left(false)      => caseTemplateTraversal.hasNotCustomField(customFieldSrv, EntityIdOrName(name))
+          }
+        case (_, caseTraversal, _, _) => caseTraversal.empty
+      }
+      .custom {
+        case (FPathElem(_, FPathElem(name, _)), value, vertex, graph, authContext) =>
+          for {
+            c <- caseTemplateSrv.get(vertex)(graph).getOrFail("CaseTemplate")
+            _ <- caseTemplateSrv.setOrCreateCustomField(c, EntityIdOrName(name), Some(value), None)(graph, authContext)
+          } yield Json.obj(s"customField.$name" -> value)
+        case (FPathElem(_, FPathEmpty), values: JsObject, vertex, graph, authContext) =>
+          for {
+            c   <- caseTemplateSrv.get(vertex)(graph).getOrFail("CaseTemplate")
+            cfv <- values.fields.toTry { case (n, v) => customFieldSrv.getOrFail(EntityIdOrName(n))(graph).map(cf => (cf, v, None)) }
+            _   <- caseTemplateSrv.updateCustomField(c, cfv)(graph, authContext)
+          } yield Json.obj("customFields" -> values)
+        case _ => Failure(BadRequestError("Invalid custom fields format"))
+      })
     .property("tasks", UMapping.jsonNative.sequence)(
       _.select(_.tasks.richTaskWithoutActionRequired.domainMap(_.toJson)).custom { //  FIXME select the correct mapping
         (_, 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
@@ -106,8 +106,7 @@ class Properties @Inject() (
       .property("customFields", UMapping.jsonNative)(_.subSelect {
         case (FPathElem(_, FPathElem(idOrName, _)), alerts) =>
           alerts
-            .customFields(EntityIdOrName(idOrName))
-            .jsonValue
+            .customFieldJsonValue(customFieldSrv, EntityIdOrName(idOrName))
         case (_, alerts) => alerts.customFields.nameJsonValue.fold.domainMap(JsObject(_))
       }
         .filter[JsValue] {
@@ -209,8 +208,7 @@ class Properties @Inject() (
       .property("customFields", UMapping.jsonNative)(_.subSelect {
         case (FPathElem(_, FPathElem(idOrName, _)), caseSteps) =>
           caseSteps
-            .customFields(EntityIdOrName(idOrName))
-            .jsonValue
+            .customFieldJsonValue(customFieldSrv, EntityIdOrName(idOrName))
         case (_, caseSteps) => caseSteps.customFields.nameJsonValue.fold.domainMap(JsObject(_))
       }
         .filter[JsValue] {
@@ -273,16 +271,34 @@ 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 <- caseTemplateSrv.getOrFail(vertex)(graph)
-            _ <- caseTemplateSrv.setOrCreateCustomField(c, name, Some(value), None)(graph, authContext)
-          } yield Json.obj(s"customField.$name" -> value)
-        case _ => Failure(BadRequestError("Invalid custom fields format"))
-      })
+        case (FPathElem(_, FPathElem(idOrName, _)), caseTemplateSteps) =>
+          caseTemplateSteps
+            .customFieldJsonValue(customFieldSrv, EntityIdOrName(idOrName))
+        case (_, caseTemplateSteps) => caseTemplateSteps.customFields.nameJsonValue.fold.domainMap(JsObject(_))
+      }
+        .filter[JsValue] {
+          case (FPathElem(_, FPathElem(name, _)), caseTemplateTraversal, _, predicate) =>
+            predicate match {
+              case Right(predicate) => caseTemplateTraversal.customFieldFilter(customFieldSrv, EntityIdOrName(name), predicate)
+              case Left(true)       => caseTemplateTraversal.hasCustomField(customFieldSrv, EntityIdOrName(name))
+              case Left(false)      => caseTemplateTraversal.hasNotCustomField(customFieldSrv, EntityIdOrName(name))
+            }
+          case (_, caseTraversal, _, _) => caseTraversal.empty
+        }
+        .custom {
+          case (FPathElem(_, FPathElem(idOrName, _)), value, vertex, graph, authContext) =>
+            for {
+              c <- caseTemplateSrv.get(vertex)(graph).getOrFail("CaseTemplate")
+              _ <- caseTemplateSrv.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   <- caseTemplateSrv.get(vertex)(graph).getOrFail("CaseTemplate")
+              cfv <- values.fields.toTry { case (n, v) => customFieldSrv.getOrFail(EntityIdOrName(n))(graph).map(cf => (cf, v, None)) }
+              _   <- caseTemplateSrv.updateCustomField(c, cfv)(graph, authContext)
+            } yield Json.obj("customFields" -> values)
+          case _ => Failure(BadRequestError("Invalid custom fields format"))
+        })
       .build
 
   lazy val organisation: PublicProperties =

diff --git a/thehive/app/org/thp/thehive/models/CustomField.scala b/thehive/app/org/thp/thehive/models/CustomField.scala
@@ -3,6 +3,9 @@ package org.thp.thehive.models
 import org.apache.tinkerpop.gremlin.structure.Edge
 import org.thp.scalligraph._
 import org.thp.scalligraph.models._
+import org.thp.scalligraph.traversal.Traversal
+import org.thp.scalligraph.traversal.Traversal.{Domain, E}
+import org.thp.scalligraph.traversal.TraversalOps._
 import play.api.libs.json._
 
 import java.util.{Date, NoSuchElementException}
@@ -91,6 +94,10 @@ sealed abstract class CustomFieldType[T] {
 
   def getJsonValue(ccf: CustomFieldValue[_]): JsValue = getValue(ccf).fold[JsValue](JsNull)(writes.writes)
 
+  def getValue[C <: CustomFieldValue[_]](traversal: Traversal.E[C]): Traversal.Domain[T]
+
+  def getJsonValue[C <: CustomFieldValue[_]](traversal: Traversal.E[C]): Traversal.Domain[JsValue] = getValue(traversal).domainMap(writes.writes)
+
   override def toString: String = name
 
   protected def setValueFailure(value: Any): Failure[Nothing] =
@@ -117,6 +124,8 @@ object CustomFieldString extends CustomFieldType[String] {
     }
 
   override def getValue(ccf: CustomFieldValue[_]): Option[String] = ccf.stringValue
+
+  override def getValue[C <: CustomFieldValue[_]](traversal: E[C]): Traversal.Domain[String] = traversal.value(_.stringValue).castDomain
 }
 
 object CustomFieldBoolean extends CustomFieldType[Boolean] {
@@ -137,6 +146,8 @@ object CustomFieldBoolean extends CustomFieldType[Boolean] {
     }
 
   override def getValue(ccf: CustomFieldValue[_]): Option[Boolean] = ccf.booleanValue
+
+  override def getValue[C <: CustomFieldValue[_]](traversal: E[C]): Domain[Boolean] = traversal.value(_.booleanValue).castDomain
 }
 
 object CustomFieldInteger extends CustomFieldType[Int] {
@@ -158,6 +169,8 @@ object CustomFieldInteger extends CustomFieldType[Int] {
     }
 
   override def getValue(ccf: CustomFieldValue[_]): Option[Int] = ccf.integerValue
+
+  override def getValue[C <: CustomFieldValue[_]](traversal: E[C]): Domain[Int] = traversal.value(_.integerValue).castDomain
 }
 
 object CustomFieldFloat extends CustomFieldType[Double] {
@@ -178,6 +191,8 @@ object CustomFieldFloat extends CustomFieldType[Double] {
     }
 
   override def getValue(ccf: CustomFieldValue[_]): Option[Double] = ccf.floatValue
+
+  override def getValue[C <: CustomFieldValue[_]](traversal: E[C]): Domain[Double] = traversal.value(_.floatValue).castDomain
 }
 
 object CustomFieldDate extends CustomFieldType[Date] {
@@ -199,6 +214,8 @@ object CustomFieldDate extends CustomFieldType[Date] {
     }
 
   override def getValue(ccf: CustomFieldValue[_]): Option[Date] = ccf.dateValue
+
+  override def getValue[C <: CustomFieldValue[_]](traversal: E[C]): Domain[Date] = traversal.value(_.dateValue).castDomain
 }
 
 @DefineIndex(IndexType.unique, "name")

diff --git a/thehive/app/org/thp/thehive/services/AlertSrv.scala b/thehive/app/org/thp/thehive/services/AlertSrv.scala
@@ -469,6 +469,14 @@ object AlertOps {
 
     def customFields: Traversal.E[AlertCustomField] = traversal.outE[AlertCustomField]
 
+    def customFieldJsonValue(customFieldSrv: CustomFieldSrv, customField: EntityIdOrName): Traversal.Domain[JsValue] =
+      customFieldSrv
+        .get(customField)(traversal.graph)
+        .value(_.`type`)
+        .headOption
+        .map(t => CustomFieldType.map(t).getJsonValue(traversal.customFields(customField)))
+        .getOrElse(traversal.empty.castDomain)
+
     def richCustomFields: Traversal[RichCustomField, JMap[String, Any], Converter[RichCustomField, JMap[String, Any]]] =
       traversal
         .outE[AlertCustomField]

diff --git a/thehive/app/org/thp/thehive/services/CaseSrv.scala b/thehive/app/org/thp/thehive/services/CaseSrv.scala
@@ -494,6 +494,14 @@ object CaseOps {
           name => customFields.filter(_.inV.v[CustomField].has(_.name, name))
         )
 
+    def customFieldJsonValue(customFieldSrv: CustomFieldSrv, customField: EntityIdOrName): Traversal.Domain[JsValue] =
+      customFieldSrv
+        .get(customField)(traversal.graph)
+        .value(_.`type`)
+        .headOption
+        .map(t => CustomFieldType.map(t).getJsonValue(traversal.customFields(customField)))
+        .getOrElse(traversal.empty.castDomain)
+
     def richCustomFields: Traversal[RichCustomField, JMap[String, Any], Converter[RichCustomField, JMap[String, Any]]] =
       customFields
         .project(_.by.by(_.inV.v[CustomField]))
+22 −22		core/src/main/scala/org/thp/scalligraph/ContextPropagatingDisptacher.scala
+2 −1		core/src/main/scala/org/thp/scalligraph/traversal/TraversalPrinter.scala
+4 −3		database/janusgraph/src/main/scala/org/thp/scalligraph/janus/IndexOps.scala
+1 −0		database/janusgraph/src/main/scala/org/thp/scalligraph/janus/JanusDatabase.scala
+6 −2		database/janusgraph/src/main/scala/org/thp/scalligraph/janus/JanusDatabaseProvider.scala