#1252 Fix format of custom fields

TheHive-Project · Mar 10, 2020 · ee3c19d · ee3c19d
1 parent 6de9456
commit ee3c19d
Show file tree

Hide file tree

Showing 6 changed files with 61 additions and 24 deletions.
diff --git a/ScalliGraph b/ScalliGraph
diff --git a/thehive/app/org/thp/thehive/controllers/v0/Properties.scala b/thehive/app/org/thp/thehive/controllers/v0/Properties.scala
@@ -97,8 +97,8 @@ class Properties @Inject() (
       .property("summary", UniMapping.string.optional)(_.field.updatable)
       .property("user", UniMapping.string)(_.field.updatable)
       .property("customFields", UniMapping.identity[JsValue])(_.subSelect {
-        case (FPathElem(_, FPathElem(name, _)), alertSteps) => alertSteps.customFields(name).jsonValue.map(_._2)
-        case (_, alertSteps)                                => alertSteps.customFields.jsonValue.fold.map(l => JsObject(l.asScala))
+        case (FPathElem(_, FPathElem(name, _)), alertSteps) => alertSteps.customFields(name).jsonValue
+        case (_, alertSteps)                                => alertSteps.customFields.nameJsonValue.fold.map(l => JsObject(l.asScala))
       }.custom {
         case (FPathElem(_, FPathElem(name, _)), value, vertex, _, graph, authContext) =>
           for {
@@ -182,12 +182,19 @@ class Properties @Inject() (
             }
           } yield Json.obj("impactStatus" -> impactStatus)
       })
-      .property("customFields", UniMapping.identity[JsValue])(_.subSelect {
-        case (FPathElem(_, FPathElem(name, _)), caseSteps) => caseSteps.customFields(name).jsonValue.map(_._2)
-        case (_, caseSteps)                                => caseSteps.customFields.jsonValue.fold.map(l => JsObject(l.asScala))
+      .property("customFields", UniMapping.jsonNative)(_.subSelect {
+        case (FPathElem(_, FPathElem(name, _)), caseSteps) =>
+          caseSteps
+            .customFields(name)
+            .value
+            .cast(UniMapping.jsonNative)
+            .get // can't fail
+
+        case (_, caseSteps) => caseSteps.customFields.nameJsonValue.fold.map(l => JsObject(l.asScala)).cast(UniMapping.jsonNative).get
       }.custom {
         case (FPathElem(_, FPathElem(name, _)), value, vertex, _, graph, authContext) =>
           for {
+//            v <- UniMapping.jsonNative.toGraphOpt(value).fold[Try[Any]](???)(Success.apply)
             c <- caseSrv.getOrFail(vertex)(graph)
             _ <- caseSrv.setOrCreateCustomField(c, name, Some(value))(graph, authContext)
           } yield Json.obj(s"customField.$name" -> value)
@@ -236,8 +243,8 @@ class Properties @Inject() (
       .property("summary", UniMapping.string.optional)(_.field.updatable)
       .property("user", UniMapping.string)(_.field.updatable)
       .property("customFields", UniMapping.identity[JsValue])(_.subSelect {
-        case (FPathElem(_, FPathElem(name, _)), caseTemplateSteps) => caseTemplateSteps.customFields(name).jsonValue.map(_._2)
-        case (_, caseTemplateSteps)                                => caseTemplateSteps.customFields.jsonValue.fold.map(l => JsObject(l.asScala))
+        case (FPathElem(_, FPathElem(name, _)), caseTemplateSteps) => caseTemplateSteps.customFields(name).jsonValue
+        case (_, caseTemplateSteps)                                => caseTemplateSteps.customFields.nameJsonValue.fold.map(l => JsObject(l.asScala))
       }.custom {
         case (FPathElem(_, FPathElem(name, _)), value, vertex, _, graph, authContext) =>
           for {

diff --git a/thehive/app/org/thp/thehive/controllers/v1/Properties.scala b/thehive/app/org/thp/thehive/controllers/v1/Properties.scala
@@ -81,8 +81,8 @@ class Properties @Inject() (
       .property("summary", UniMapping.string.optional)(_.field.updatable)
       .property("user", UniMapping.string)(_.field.updatable)
       .property("customFields", UniMapping.identity[JsValue])(_.subSelect {
-        case (FPathElem(_, FPathElem(name, _)), alertSteps) => alertSteps.customFields(name).jsonValue.map(_._2)
-        case (_, alertSteps)                                => alertSteps.customFields.jsonValue.fold.map(l => JsObject(l.asScala))
+        case (FPathElem(_, FPathElem(name, _)), alertSteps) => alertSteps.customFields(name).jsonValue
+        case (_, alertSteps)                                => alertSteps.customFields.nameJsonValue.fold.map(l => JsObject(l.asScala))
       }.custom {
         case (FPathElem(_, FPathElem(name, _)), value, vertex, _, graph, authContext) =>
           for {
@@ -162,8 +162,8 @@ class Properties @Inject() (
       .property("summary", UniMapping.string.optional)(_.field.updatable)
       .property("user", UniMapping.string)(_.field.updatable)
       .property("customFields", UniMapping.identity[JsValue])(_.subSelect {
-        case (FPathElem(_, FPathElem(name, _)), alertSteps) => alertSteps.customFields(name).jsonValue.map(_._2)
-        case (_, alertSteps)                                => alertSteps.customFields.jsonValue.fold.map(l => JsObject(l.asScala))
+        case (FPathElem(_, FPathElem(name, _)), alertSteps) => alertSteps.customFields(name).jsonValue
+        case (_, alertSteps)                                => alertSteps.customFields.nameJsonValue.fold.map(l => JsObject(l.asScala))
       }.custom {
         case (FPathElem(_, FPathElem(name, _)), value, vertex, _, graph, authContext) =>
           for {

diff --git a/thehive/app/org/thp/thehive/services/CustomFieldSrv.scala b/thehive/app/org/thp/thehive/services/CustomFieldSrv.scala
@@ -5,7 +5,7 @@ import java.util.{Map => JMap}
 import scala.collection.JavaConverters._
 import scala.util.Try
 
-import play.api.libs.json.{JsObject, JsValue}
+import play.api.libs.json.{JsNull, JsObject, JsValue}
 
 import gremlin.scala._
 import javax.inject.{Inject, Singleton}
@@ -101,23 +101,53 @@ class CustomFieldValueSteps(raw: GremlinScala[Edge])(implicit db: Database, grap
       .map(_ => ())
   }
 
-  def jsonValue: Traversal[(String, JsValue), (String, JsValue)] = {
+  private def edgeNameType: GremlinScala[(Edge, JMap[AnyRef, AnyRef])] = {
     val customFieldValueLabel = StepLabel[Edge]()
     val typeLabel             = StepLabel[JMap[AnyRef, AnyRef]]()
+    raw
+      .asInstanceOf[GremlinScala.Aux[Edge, HNil]]
+      .as(customFieldValueLabel)
+      .inV()
+      .valueMap("name", "type")
+      .as(typeLabel)
+      .select()
+  }
+
+  def nameJsonValue: Traversal[(String, JsValue), (String, JsValue)] =
     Traversal(
-      raw
-        .asInstanceOf[GremlinScala.Aux[Edge, HNil]]
-        .as(customFieldValueLabel)
-        .inV()
-        .valueMap("name", "type")
-        .as(typeLabel)
-        .select()
+      edgeNameType
         .map {
           case (edge, ValueMap(map)) =>
             map.get[String]("name") -> CustomFieldType.get(map.get[String]("type")).getJsonValue(new CustomFieldValueEdge(db, edge))
         }
     )
-  }
+
+  def jsonValue: Traversal[JsValue, JsValue] =
+    Traversal(
+      edgeNameType
+        .map {
+          case (edge, ValueMap(map)) =>
+            CustomFieldType.get(map.get[String]("type")).getJsonValue(new CustomFieldValueEdge(db, edge))
+        }
+    )
+
+  def nameValue: Traversal[(String, Option[Any]), (String, Option[Any])] =
+    Traversal(
+      edgeNameType
+        .map {
+          case (edge, ValueMap(map)) =>
+            map.get[String]("name") -> CustomFieldType.get(map.get[String]("type")).getValue(new CustomFieldValueEdge(db, edge))
+        }
+    )
+
+  def value: Traversal[Any, Any] =
+    Traversal(
+      edgeNameType
+        .map {
+          case (edge, ValueMap(map)) =>
+            CustomFieldType.get(map.get[String]("type")).getValue(new CustomFieldValueEdge(db, edge)).getOrElse(JsNull)
+        }
+    )
 
   def richCustomField: Traversal[RichCustomField, RichCustomField] = {
     val customFieldValueLabel = StepLabel[Edge]()

diff --git a/thehive/test/org/thp/thehive/services/AlertSrvTest.scala b/thehive/test/org/thp/thehive/services/AlertSrvTest.scala
@@ -138,7 +138,7 @@ class AlertSrvTest extends PlaySpecification with TestAppBuilder {
       } must beSuccessfulTry
 
       app[Database].roTransaction { implicit graph =>
-        app[AlertSrv].get("testType;testSource;ref1").customFields("string1").jsonValue.headOption()
+        app[AlertSrv].get("testType;testSource;ref1").customFields("string1").nameJsonValue.headOption()
       } must beSome("string1" -> JsString("sad"))
     }
 

diff --git a/thehive/test/org/thp/thehive/services/CaseTemplateSrvTest.scala b/thehive/test/org/thp/thehive/services/CaseTemplateSrvTest.scala
@@ -110,7 +110,7 @@ class CaseTemplateSrvTest extends PlaySpecification with TestAppBuilder {
 
       val expected: Seq[(String, JsValue)] = Seq("string1" -> JsString("hate"), "boolean1" -> JsTrue, "integer1" -> JsNumber(1))
       app[Database].roTransaction { implicit graph =>
-        app[CaseTemplateSrv].get("spam").customFields.jsonValue.toList
+        app[CaseTemplateSrv].get("spam").customFields.nameJsonValue.toList
       } must contain(exactly(expected: _*))
     }
   }