From b36f8e8443cbbe726ff687dd9ae11be2f14af3b6 Mon Sep 17 00:00:00 2001 From: To-om Date: Fri, 5 Mar 2021 15:45:29 +0100 Subject: [PATCH] #1670 Add extra data on tags to see usage --- .../scala/org/thp/thehive/dto/v1/Tag.scala | 5 ++- .../thehive/controllers/v1/Conversion.scala | 18 ++++++++ .../thp/thehive/controllers/v1/TagCtrl.scala | 10 +++-- .../thehive/controllers/v1/TagRenderer.scala | 42 +++++++++++++++++++ .../app/org/thp/thehive/services/TagSrv.scala | 11 ++++- 5 files changed, 80 insertions(+), 6 deletions(-) create mode 100644 thehive/app/org/thp/thehive/controllers/v1/TagRenderer.scala diff --git a/dto/src/main/scala/org/thp/thehive/dto/v1/Tag.scala b/dto/src/main/scala/org/thp/thehive/dto/v1/Tag.scala index c2291cd4a3..a32227a520 100644 --- a/dto/src/main/scala/org/thp/thehive/dto/v1/Tag.scala +++ b/dto/src/main/scala/org/thp/thehive/dto/v1/Tag.scala @@ -1,6 +1,6 @@ package org.thp.thehive.dto.v1 -import play.api.libs.json.{Json, OFormat} +import play.api.libs.json.{JsObject, Json, OFormat} import java.util.Date @@ -15,7 +15,8 @@ case class OutputTag( predicate: String, value: Option[String], description: Option[String], - colour: String + colour: String, + extraData: JsObject ) object OutputTag { diff --git a/thehive/app/org/thp/thehive/controllers/v1/Conversion.scala b/thehive/app/org/thp/thehive/controllers/v1/Conversion.scala index 52b80a53fa..0997186bbb 100644 --- a/thehive/app/org/thp/thehive/controllers/v1/Conversion.scala +++ b/thehive/app/org/thp/thehive/controllers/v1/Conversion.scala @@ -311,9 +311,27 @@ object Conversion { .withFieldConst(_._createdBy, tag._createdBy) .withFieldConst(_._type, "Tag") .withFieldComputed(_.namespace, t => if (t.isFreeTag) "_freetags_" else t.namespace) + .withFieldConst(_.extraData, JsObject.empty) .transform ) + implicit val tagWithStatsOutput: Renderer.Aux[(Tag with Entity, JsObject), OutputTag] = + Renderer.toJson[(Tag with Entity, JsObject), OutputTag] { + case (tag, stats) => + tag + .asInstanceOf[Tag] + .into[OutputTag] + .withFieldConst(_._id, tag._id.toString) + .withFieldConst(_._updatedAt, tag._updatedAt) + .withFieldConst(_._updatedBy, tag._updatedBy) + .withFieldConst(_._createdAt, tag._createdAt) + .withFieldConst(_._createdBy, tag._createdBy) + .withFieldConst(_._type, "Tag") + .withFieldComputed(_.namespace, t => if (t.isFreeTag) "_freetags_" else t.namespace) + .withFieldConst(_.extraData, stats) + .transform + } + implicit class InputUserOps(inputUser: InputUser) { def toUser: User = diff --git a/thehive/app/org/thp/thehive/controllers/v1/TagCtrl.scala b/thehive/app/org/thp/thehive/controllers/v1/TagCtrl.scala index f474b03b07..4c83a40857 100644 --- a/thehive/app/org/thp/thehive/controllers/v1/TagCtrl.scala +++ b/thehive/app/org/thp/thehive/controllers/v1/TagCtrl.scala @@ -23,16 +23,20 @@ class TagCtrl @Inject() ( entrypoint: Entrypoint, db: Database, tagSrv: TagSrv, - organisationSrv: OrganisationSrv, + val organisationSrv: OrganisationSrv, properties: Properties -) extends QueryableCtrl { +) extends QueryableCtrl + with TagRenderer { override val entityName: String = "Tag" override val publicProperties: PublicProperties = properties.tag override val initialQuery: Query = Query.init[Traversal.V[Tag]]("listTag", (graph, authContext) => tagSrv.startTraversal(graph).visible(authContext)) override val pageQuery: ParamQuery[OutputParam] = Query.withParam[OutputParam, Traversal.V[Tag], IteratorOutput]( "page", - (range, tagSteps, _) => tagSteps.page(range.from, range.to, withTotal = true) + (params, tagSteps, authContext) => + tagSteps.richPage(params.from, params.to, params.extraData.contains("total"))( + _.withCustomRenderer(tagStatsRenderer(params.extraData - "total")(authContext)) + ) ) override val outputQuery: Query = Query.output[Tag with Entity] override val getQuery: ParamQuery[EntityIdOrName] = Query.initWithParam[EntityIdOrName, Traversal.V[Tag]]( diff --git a/thehive/app/org/thp/thehive/controllers/v1/TagRenderer.scala b/thehive/app/org/thp/thehive/controllers/v1/TagRenderer.scala new file mode 100644 index 0000000000..bc9fe614e5 --- /dev/null +++ b/thehive/app/org/thp/thehive/controllers/v1/TagRenderer.scala @@ -0,0 +1,42 @@ +package org.thp.thehive.controllers.v1 + +import org.thp.scalligraph.auth.AuthContext +import org.thp.scalligraph.traversal.TraversalOps._ +import org.thp.scalligraph.traversal.{Converter, Traversal} +import org.thp.thehive.models.Tag +import org.thp.thehive.services.TagOps._ +import play.api.libs.json._ + +import java.util.{Map => JMap} + +trait TagRenderer extends BaseRenderer[Tag] { + + def usageStats(implicit + authContext: AuthContext + ): Traversal.V[Tag] => Traversal[JsObject, JMap[String, Any], Converter[JsObject, JMap[String, Any]]] = + _.project( + _.by(_.`case`.count) + .by(_.alert.count) + .by(_.observable.count) + ).domainMap { + case (caseCount, alertCount, observableCount) => + Json.obj( + "case" -> caseCount, + "alert" -> alertCount, + "observable" -> observableCount + ) + } + + def tagStatsRenderer(extraData: Set[String])(implicit + authContext: AuthContext + ): Traversal.V[Tag] => JsTraversal = { implicit traversal => + baseRenderer( + extraData, + traversal, + { + case (f, "usage") => addData("usage", f)(usageStats) + case (f, _) => f + } + ) + } +} diff --git a/thehive/app/org/thp/thehive/services/TagSrv.scala b/thehive/app/org/thp/thehive/services/TagSrv.scala index 60b3390088..823da0f84e 100644 --- a/thehive/app/org/thp/thehive/services/TagSrv.scala +++ b/thehive/app/org/thp/thehive/services/TagSrv.scala @@ -1,5 +1,6 @@ package org.thp.thehive.services +import java.util.{Map => JMap} import akka.actor.ActorRef import org.apache.tinkerpop.gremlin.process.traversal.TextP import org.apache.tinkerpop.gremlin.structure.Vertex @@ -109,10 +110,13 @@ object TagOps { def displayName: Traversal[String, Vertex, Converter[String, Vertex]] = traversal.domainMap(_.toString) def fromCase: Traversal.V[Tag] = traversal.filter(_.in[CaseTag]) + def `case`: Traversal.V[Case] = traversal.in[CaseTag].v[Case] - def fromObservable: Traversal.V[Tag] = traversal.filter(_.in[ObservableTag]) + def fromObservable: Traversal.V[Tag] = traversal.filter(_.in[ObservableTag]) + def observable: Traversal.V[Observable] = traversal.in[ObservableTag].v[Observable] def fromAlert: Traversal.V[Tag] = traversal.filter(_.in[AlertTag]) + def alert: Traversal.V[Alert] = traversal.in[AlertTag].v[Alert] def freetags(organisationSrv: OrganisationSrv)(implicit authContext: AuthContext): Traversal.V[Tag] = { val freeTagNamespace: String = s"_freetags_${organisationSrv.currentId(traversal.graph, authContext).value}" @@ -140,6 +144,11 @@ object TagOps { def visible(implicit authContext: AuthContext): Traversal.V[Tag] = traversal.filter(_.organisation.current) + + def withCustomRenderer[D, G, C <: Converter[D, G]]( + entityRenderer: Traversal.V[Tag] => Traversal[D, G, C] + ): Traversal[(Tag with Entity, D), JMap[String, Any], Converter[(Tag with Entity, D), JMap[String, Any]]] = + traversal.project(_.by.by(entityRenderer)) } }