From 51405c96fd7ac78cbdb2d47e750c8c4807e16335 Mon Sep 17 00:00:00 2001 From: To-om Date: Tue, 10 Nov 2020 16:06:26 +0100 Subject: [PATCH] #1636 Fix customField filter in alert --- ScalliGraph | 2 +- .../thehive/controllers/v0/AlertCtrl.scala | 2 +- .../thehive/controllers/v1/Properties.scala | 61 +++++++++++++++---- .../org/thp/thehive/services/AlertSrv.scala | 16 +++-- .../thp/thehive/services/AlertSrvTest.scala | 4 +- 5 files changed, 64 insertions(+), 21 deletions(-) diff --git a/ScalliGraph b/ScalliGraph index f3e9b712c1..0dc00d560b 160000 --- a/ScalliGraph +++ b/ScalliGraph @@ -1 +1 @@ -Subproject commit f3e9b712c15bf23f715968bbca4d00d39f6ff1db +Subproject commit 0dc00d560b7c48f5f7a781cd4c9861a64f56cd23 diff --git a/thehive/app/org/thp/thehive/controllers/v0/AlertCtrl.scala b/thehive/app/org/thp/thehive/controllers/v0/AlertCtrl.scala index 8c71bea2f8..a3187ae1f3 100644 --- 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 index 0a3a7c8619..fa40aae20b 100644 --- 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 index 5772d21d5f..b9877b41e2 100644 --- 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 index 5db4963a9a..8d6a0e2bcb 100644 --- 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")) }