#1410 WIP: add case stats

TheHive-Project · Jun 28, 2020 · 19a1452 · 19a1452
1 parent 097f3aa
commit 19a1452
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 21 deletions.
diff --git a/thehive/app/org/thp/thehive/controllers/v0/CaseRenderer.scala b/thehive/app/org/thp/thehive/controllers/v0/CaseRenderer.scala
@@ -3,7 +3,6 @@ package org.thp.thehive.controllers.v0
 import java.lang.{Long => JLong}
 
 import gremlin.scala.{__, By, Graph, GremlinScala, Key, Vertex}
-import javax.inject.Named
 import org.thp.scalligraph.auth.AuthContext
 import org.thp.scalligraph.models.Database
 import org.thp.scalligraph.services._
@@ -19,14 +18,14 @@ trait CaseRenderer {
 
   def observableStats(
       shareTraversal: GremlinScala[Vertex]
-  )(implicit @Named("with-thehive-schema") db: Database, graph: Graph): GremlinScala[JsObject] =
+  )(implicit db: Database, graph: Graph): GremlinScala[JsObject] =
     new ShareSteps(shareTraversal)
       .observables
       .count
       .map(count => Json.obj("count" -> count))
       .raw
 
-  def taskStats(shareTraversal: GremlinScala[Vertex])(implicit @Named("with-thehive-schema") db: Database, graph: Graph): GremlinScala[JsObject] =
+  def taskStats(shareTraversal: GremlinScala[Vertex])(implicit db: Database, graph: Graph): GremlinScala[JsObject] =
     new ShareSteps(shareTraversal)
       .tasks
       .active
@@ -60,23 +59,23 @@ trait CaseRenderer {
 
   def sharedWithStats(
       caseTraversal: GremlinScala[Vertex]
-  )(implicit @Named("with-thehive-schema") db: Database, graph: Graph): GremlinScala[Seq[String]] =
+  )(implicit db: Database, graph: Graph): GremlinScala[Seq[String]] =
     new CaseSteps(caseTraversal).organisations.name.fold.map(_.asScala.toSeq).raw
 
-  def originStats(caseTraversal: GremlinScala[Vertex])(implicit @Named("with-thehive-schema") db: Database, graph: Graph): GremlinScala[String] =
+  def originStats(caseTraversal: GremlinScala[Vertex])(implicit db: Database, graph: Graph): GremlinScala[String] =
     new CaseSteps(caseTraversal).origin.name.raw
 
-  def shareCountStats(caseTraversal: GremlinScala[Vertex])(implicit @Named("with-thehive-schema") db: Database, graph: Graph): GremlinScala[JLong] =
+  def shareCountStats(caseTraversal: GremlinScala[Vertex])(implicit db: Database, graph: Graph): GremlinScala[JLong] =
     new CaseSteps(caseTraversal).organisations.count.raw
 
   def isOwnerStats(
       caseTraversal: GremlinScala[Vertex]
-  )(implicit @Named("with-thehive-schema") db: Database, graph: Graph, authContext: AuthContext): GremlinScala[Boolean] =
+  )(implicit db: Database, graph: Graph, authContext: AuthContext): GremlinScala[Boolean] =
     new CaseSteps(caseTraversal).origin.name.map(_ == authContext.organisation).raw
 
   def caseStatsRenderer(
       implicit authContext: AuthContext,
-      @Named("with-thehive-schema") db: Database,
+      db: Database,
       graph: Graph
   ): CaseSteps => Traversal[JsObject, JsObject] =
     _.project(

diff --git a/thehive/app/org/thp/thehive/controllers/v1/CaseCtrl.scala b/thehive/app/org/thp/thehive/controllers/v1/CaseCtrl.scala
@@ -11,6 +11,7 @@ import org.thp.thehive.controllers.v1.Conversion._
 import org.thp.thehive.dto.v1.{InputCase, InputTask}
 import org.thp.thehive.models.{Permissions, RichCase, User}
 import org.thp.thehive.services._
+import play.api.libs.json.JsObject
 import play.api.mvc.{Action, AnyContent, Results}
 
 import scala.util.{Success, Try}
@@ -25,7 +26,8 @@ class CaseCtrl @Inject() (
     userSrv: UserSrv,
     tagSrv: TagSrv,
     organisationSrv: OrganisationSrv
-) extends QueryableCtrl {
+) extends QueryableCtrl
+    with CaseRenderer {
 
   override val entityName: String                           = "case"
   override val publicProperties: List[PublicProperty[_, _]] = properties.`case` ::: metaProperties[CaseSteps]
@@ -36,18 +38,13 @@ class CaseCtrl @Inject() (
     FieldsParser[IdOrName],
     (param, graph, authContext) => caseSrv.get(param.idOrName)(graph).visible(authContext)
   )
-  override val pageQuery: ParamQuery[OutputParam] = Query.withParam[OutputParam, CaseSteps, PagedResult[RichCase]](
+  override val pageQuery: ParamQuery[OutputParam] = Query.withParam[OutputParam, CaseSteps, PagedResult[(RichCase, JsObject)]](
     "page",
     FieldsParser[OutputParam], {
-      case (OutputParam(from, to, withStats), caseSteps, authContext) =>
+      case (OutputParam(from, to, extraData), caseSteps, authContext) =>
         caseSteps
-          .richPage(from, to, withTotal = true) { c =>
-            c.richCase(authContext)
-//            case c if withStats =>
-////              c.richCaseWithCustomRenderer(caseStatsRenderer(authContext, db, caseSteps.graph))
-//              c.richCase.map(_ -> JsObject.empty) // TODO add stats
-//            case c =>
-//              c.richCase.map(_ -> JsObject.empty)
+          .richPage(from, to, extraData.contains("total")) { c =>
+            c.richCaseWithCustomRenderer(caseStatsRenderer(extraData - "total")(authContext, db, caseSteps.graph))(authContext)
           }
     }
   )

diff --git a/thehive/app/org/thp/thehive/controllers/v1/TheHiveQueryExecutor.scala b/thehive/app/org/thp/thehive/controllers/v1/TheHiveQueryExecutor.scala
@@ -5,16 +5,16 @@ import org.thp.scalligraph.controllers.{FObject, FieldsParser}
 import org.thp.scalligraph.models.Database
 import org.thp.scalligraph.query._
 
-case class OutputParam(from: Long, to: Long, withStats: Boolean)
+case class OutputParam(from: Long, to: Long, extraData: Set[String])
 
 object OutputParam {
   implicit val parser: FieldsParser[OutputParam] = FieldsParser[OutputParam]("OutputParam") {
     case (_, field: FObject) =>
       for {
         from      <- FieldsParser.long.on("from")(field)
         to        <- FieldsParser.long.on("to")(field)
-        withStats <- FieldsParser.boolean.optional.on("withStats")(field)
-      } yield OutputParam(from, to, withStats.getOrElse(false))
+        extraData <- FieldsParser.string.set.on("extraData")(field)
+      } yield OutputParam(from, to, extraData)
   }
 }