diff --git a/dto/src/main/scala/org/thp/thehive/dto/v1/Procedure.scala b/dto/src/main/scala/org/thp/thehive/dto/v1/Procedure.scala index fc57cd9cb7..a64e4d84a9 100644 --- a/dto/src/main/scala/org/thp/thehive/dto/v1/Procedure.scala +++ b/dto/src/main/scala/org/thp/thehive/dto/v1/Procedure.scala @@ -1,6 +1,6 @@ package org.thp.thehive.dto.v1 -import play.api.libs.json.{Format, Json, Reads, Writes} +import play.api.libs.json.{Format, JsObject, Json, Reads, Writes} import java.util.Date @@ -37,7 +37,8 @@ case class OutputProcedure( _updatedBy: Option[String], description: String, occurence: Date, - patternId: String + patternId: String, + extraData: JsObject ) object OutputProcedure { diff --git a/thehive/app/org/thp/thehive/controllers/v1/CaseCtrl.scala b/thehive/app/org/thp/thehive/controllers/v1/CaseCtrl.scala index e52b9b8b75..182cdbc688 100644 --- a/thehive/app/org/thp/thehive/controllers/v1/CaseCtrl.scala +++ b/thehive/app/org/thp/thehive/controllers/v1/CaseCtrl.scala @@ -60,7 +60,8 @@ class CaseCtrl @Inject() ( Query[Traversal.V[Case], Traversal.V[User]]("assignableUsers", (caseSteps, authContext) => caseSteps.assignableUsers(authContext)), Query[Traversal.V[Case], Traversal.V[Organisation]]("organisations", (caseSteps, authContext) => caseSteps.organisations.visible(authContext)), Query[Traversal.V[Case], Traversal.V[Alert]]("alerts", (caseSteps, authContext) => caseSteps.alert.visible(authContext)), - Query[Traversal.V[Case], Traversal.V[Share]]("shares", (caseSteps, authContext) => caseSteps.shares.visible(authContext)) + Query[Traversal.V[Case], Traversal.V[Share]]("shares", (caseSteps, authContext) => caseSteps.shares.visible(authContext)), + Query[Traversal.V[Case], Traversal.V[Procedure]]("procedures", (caseSteps, _) => caseSteps.procedure) ) def create: Action[AnyContent] = diff --git a/thehive/app/org/thp/thehive/controllers/v1/Conversion.scala b/thehive/app/org/thp/thehive/controllers/v1/Conversion.scala index de76eef468..177f77ecb7 100644 --- a/thehive/app/org/thp/thehive/controllers/v1/Conversion.scala +++ b/thehive/app/org/thp/thehive/controllers/v1/Conversion.scala @@ -546,7 +546,19 @@ object Conversion { _.into[OutputProcedure] .withFieldComputed(_._id, _._id.toString) .withFieldComputed(_.patternId, _.pattern.patternId) + .withFieldConst(_.extraData, JsObject.empty) .transform ) + implicit val richProcedureWithStatsRenderer: Renderer.Aux[(RichProcedure, JsObject), OutputProcedure] = + Renderer.toJson[(RichProcedure, JsObject), OutputProcedure] { procedureWithExtraData => + procedureWithExtraData + ._1 + .into[OutputProcedure] + .withFieldComputed(_._id, _._id.toString) + .withFieldComputed(_.patternId, _.pattern.patternId) + .withFieldConst(_.extraData, procedureWithExtraData._2) + .transform + } + } diff --git a/thehive/app/org/thp/thehive/controllers/v1/ProcedureCtrl.scala b/thehive/app/org/thp/thehive/controllers/v1/ProcedureCtrl.scala index bd4e18219b..7225aa21e2 100644 --- a/thehive/app/org/thp/thehive/controllers/v1/ProcedureCtrl.scala +++ b/thehive/app/org/thp/thehive/controllers/v1/ProcedureCtrl.scala @@ -21,7 +21,8 @@ class ProcedureCtrl @Inject() ( properties: Properties, procedureSrv: ProcedureSrv, @Named("with-thehive-schema") implicit val db: Database -) extends QueryableCtrl { +) extends QueryableCtrl + with ProcedureRenderer { override val entityName: String = "procedure" override val publicProperties: PublicProperties = properties.procedure override val initialQuery: Query = Query.init[Traversal.V[Procedure]]( @@ -33,7 +34,10 @@ class ProcedureCtrl @Inject() ( override val pageQuery: ParamQuery[OutputParam] = Query.withParam[OutputParam, Traversal.V[Procedure], IteratorOutput]( "page", FieldsParser[OutputParam], - (range, procedureSteps, _) => procedureSteps.richPage(range.from, range.to, range.extraData.contains("total"))(_.richProcedure) + (range, procedureSteps, _) => + procedureSteps.richPage(range.from, range.to, range.extraData.contains("total"))( + _.richProcedureWithCustomRenderer(procedureStatsRenderer(range.extraData - "total")) + ) ) override val outputQuery: Query = Query.output[RichProcedure, Traversal.V[Procedure]](_.richProcedure) override val getQuery: ParamQuery[EntityIdOrName] = Query.initWithParam[EntityIdOrName, Traversal.V[Procedure]]( diff --git a/thehive/app/org/thp/thehive/controllers/v1/ProcedureRenderer.scala b/thehive/app/org/thp/thehive/controllers/v1/ProcedureRenderer.scala new file mode 100644 index 0000000000..943488c375 --- /dev/null +++ b/thehive/app/org/thp/thehive/controllers/v1/ProcedureRenderer.scala @@ -0,0 +1,27 @@ +package org.thp.thehive.controllers.v1 + +import org.thp.scalligraph.traversal.{Converter, Traversal} +import org.thp.thehive.controllers.v1.Conversion._ +import org.thp.thehive.models.Procedure +import org.thp.thehive.services.PatternOps._ +import org.thp.thehive.services.ProcedureOps._ +import play.api.libs.json.JsValue + +import java.util.{Map => JMap} + +trait ProcedureRenderer extends BaseRenderer[Procedure] { + def patternStats: Traversal.V[Procedure] => Traversal[JsValue, JMap[String, Any], Converter[JsValue, JMap[String, Any]]] = + _.pattern.richPattern.domainMap(_.toJson) + + def procedureStatsRenderer(extraData: Set[String]): Traversal.V[Procedure] => JsTraversal = { implicit traversal => + baseRenderer( + extraData, + traversal, + { + case (f, "pattern") => addData("pattern", f)(patternStats) + case (f, _) => f + } + ) + } + +} diff --git a/thehive/app/org/thp/thehive/services/ProcedureSrv.scala b/thehive/app/org/thp/thehive/services/ProcedureSrv.scala index 15e81cdde0..d5860c573c 100644 --- a/thehive/app/org/thp/thehive/services/ProcedureSrv.scala +++ b/thehive/app/org/thp/thehive/services/ProcedureSrv.scala @@ -80,5 +80,18 @@ object ProcedureOps { .by(_.pattern) ) .domainMap { case (procedure, pattern) => RichProcedure(procedure, pattern) } + + def richProcedureWithCustomRenderer[D, G, C <: Converter[D, G]]( + entityRenderer: Traversal.V[Procedure] => Traversal[D, G, C] + ): Traversal[(RichProcedure, D), JMap[String, Any], Converter[(RichProcedure, D), JMap[String, Any]]] = + traversal + .project( + _.by + .by(_.pattern) + .by(entityRenderer) + ) + .domainMap { + case (procedure, pattern, renderedEntity) => (RichProcedure(procedure, pattern), renderedEntity) + } } }