#954 Add total in custom field use count

TheHive-Project · Apr 6, 2020 · fb941c0 · fb941c0
1 parent 6aebece
commit fb941c0
Show file tree

Hide file tree

Showing 5 changed files with 33 additions and 29 deletions.
diff --git a/thehive-backend/app/controllers/AlertCtrl.scala b/thehive-backend/app/controllers/AlertCtrl.scala
@@ -128,8 +128,9 @@ class AlertCtrl @Inject()(
   @Timed
   def bulkDelete(): Action[Fields] = authenticated(Roles.admin).async(fieldsBodyParser) { implicit request ⇒
     request.body.getStrings("ids").fold(Future.successful(NoContent)) { ids ⇒
-      Future.traverse(ids)(alertSrv.delete(_, request.body.getBoolean("force").getOrElse(false)))
-        .map(_ => NoContent)
+      Future
+        .traverse(ids)(alertSrv.delete(_, request.body.getBoolean("force").getOrElse(false)))
+        .map(_ ⇒ NoContent)
     }
   }
 
@@ -180,7 +181,9 @@ class AlertCtrl @Inject()(
   def createCase(id: String): Action[Fields] = authenticated(Roles.write).async(fieldsBodyParser) { implicit request ⇒
     for {
       alert ← alertSrv.get(id)
-      customCaseTemplate = request.body.getString("caseTemplate")
+      customCaseTemplate = request
+        .body
+        .getString("caseTemplate")
         .orElse(alert.caseTemplate())
       caze ← alertSrv.createCase(alert, customCaseTemplate)
     } yield renderer.toOutput(CREATED, caze)

diff --git a/thehive-backend/app/controllers/CustomFieldsCtrl.scala b/thehive-backend/app/controllers/CustomFieldsCtrl.scala
@@ -1,15 +1,15 @@
 package controllers
 
-import scala.concurrent.{ ExecutionContext, Future }
+import scala.concurrent.{ExecutionContext, Future}
 
 import play.api.http.Status
-import play.api.libs.json.{ JsNumber, JsObject }
-import play.api.mvc.{ AbstractController, Action, AnyContent, ControllerComponents }
+import play.api.libs.json.{JsNumber, JsObject, Json}
+import play.api.mvc.{AbstractController, Action, AnyContent, ControllerComponents}
 
 import akka.stream.Materializer
 import akka.stream.scaladsl.Sink
-import com.sksamuel.elastic4s.http.ElasticDsl.{ search, termsAggregation }
-import javax.inject.{ Inject, Singleton }
+import com.sksamuel.elastic4s.http.ElasticDsl.{search, termsAggregation}
+import javax.inject.{Inject, Singleton}
 import models.Roles
 
 import org.elastic4play.NotFoundError
@@ -45,12 +45,12 @@ class CustomFieldsCtrl @Inject()(
             dbfind(
               indexName ⇒ search(indexName).query(filter.query).aggregations(termsAggregation("t").field("relations"))
             ).map { searchResponse ⇒
-                val result = JsObject(searchResponse.aggregations.terms("t").buckets.map { b ⇒
-                  b.key → JsNumber(b.docCount)
-                })
-                Ok(result)
+                val buckets = searchResponse.aggregations.terms("t").buckets
+                val total   = buckets.map(_.docCount).sum
+                val result  = buckets.map(b ⇒ b.key → JsNumber(b.docCount)) :+ ("total" → JsNumber(total))
+                Ok(JsObject(result))
               }
-              .recover { case _ ⇒ Ok(JsObject.empty) }
+              .recover { case _ ⇒ Ok(Json.obj("total" → 0)) }
           }
       }
 }
diff --git a/thehive-backend/app/controllers/StatusCtrl.scala b/thehive-backend/app/controllers/StatusCtrl.scala
@@ -33,8 +33,8 @@ class StatusCtrl @Inject()(
 ) extends AbstractController(components) {
 
   private[controllers] def getVersion(c: Class[_]) = Option(c.getPackage.getImplementationVersion).getOrElse("SNAPSHOT")
-  private var clusterStatusName: String = "Init"
-  val checkStatusInterval: FiniteDuration = configuration.getOptional[FiniteDuration]("statusCheckInterval").getOrElse(1.minute)
+  private var clusterStatusName: String            = "Init"
+  val checkStatusInterval: FiniteDuration          = configuration.getOptional[FiniteDuration]("statusCheckInterval").getOrElse(1.minute)
   private def updateStatus(): Unit = {
     clusterStatusName = Try(dbIndex.clusterStatusName).getOrElse("ERROR")
     system.scheduler.scheduleOnce(checkStatusInterval)(updateStatus())

diff --git a/thehive-backend/app/controllers/StreamCtrl.scala b/thehive-backend/app/controllers/StreamCtrl.scala
@@ -23,7 +23,7 @@ import services.StreamActor
 import services.StreamActor.StreamMessages
 
 import org.elastic4play.controllers._
-import org.elastic4play.services.{ AuxSrv, EventSrv, MigrationSrv, UserSrv }
+import org.elastic4play.services.{AuxSrv, EventSrv, MigrationSrv, UserSrv}
 import org.elastic4play.Timed
 
 @Singleton
@@ -99,9 +99,10 @@ class StreamCtrl(
       Future.successful(BadRequest("Invalid stream id"))
     } else {
       val futureStatus = authenticated.expirationStatus(request) match {
-        case ExpirationError if !migrationSrv.isMigrating ⇒ userSrv.getInitialUser(request).recoverWith { case _ => authenticated.getFromApiKey(request)}.map(_ ⇒ OK)
-        case _: ExpirationWarning                         ⇒ Future.successful(220)
-        case _                                            ⇒ Future.successful(OK)
+        case ExpirationError if !migrationSrv.isMigrating ⇒
+          userSrv.getInitialUser(request).recoverWith { case _ ⇒ authenticated.getFromApiKey(request) }.map(_ ⇒ OK)
+        case _: ExpirationWarning ⇒ Future.successful(220)
+        case _                    ⇒ Future.successful(OK)
       }
 
       // Check if stream actor exists

diff --git a/thehive-backend/app/services/AlertSrv.scala b/thehive-backend/app/services/AlertSrv.scala
@@ -3,27 +3,27 @@ package services
 import java.nio.file.Files
 
 import scala.collection.immutable
-import scala.concurrent.{ ExecutionContext, Future }
+import scala.concurrent.{ExecutionContext, Future}
 import scala.util.matching.Regex
-import scala.util.{ Failure, Success, Try }
+import scala.util.{Failure, Success, Try}
 
 import play.api.libs.json._
-import play.api.{ Configuration, Logger }
+import play.api.{Configuration, Logger}
 
 import akka.NotUsed
 import akka.stream.Materializer
-import akka.stream.scaladsl.{ Sink, Source }
+import akka.stream.scaladsl.{Sink, Source}
 import connectors.ConnectorRouter
-import javax.inject.{ Inject, Singleton }
+import javax.inject.{Inject, Singleton}
 import models._
 
-import org.elastic4play.controllers.{ Fields, FileInputValue }
+import org.elastic4play.controllers.{Fields, FileInputValue}
 import org.elastic4play.database.ModifyConfig
 import org.elastic4play.services.JsonFormat.attachmentFormat
-import org.elastic4play.services.QueryDSL.{ groupByField, parent, selectCount, withId }
+import org.elastic4play.services.QueryDSL.{groupByField, parent, selectCount, withId}
 import org.elastic4play.services._
 import org.elastic4play.utils.Collection
-import org.elastic4play.{ ConflictError, InternalError }
+import org.elastic4play.{ConflictError, InternalError}
 
 trait AlertTransformer {
   def createCase(alert: Alert, customCaseTemplate: Option[String])(implicit authContext: AuthContext): Future[Case]
@@ -312,7 +312,7 @@ class AlertSrv(
       .create(caze, artifactsFields)
       .flatMap { artifacts ⇒
         Future.traverse(artifacts) {
-          case Success(_) => Future.successful(())
+          case Success(_) ⇒ Future.successful(())
           case Failure(ConflictError(_, attributes)) ⇒ // if it already exists, add tags from alert
             import org.elastic4play.services.QueryDSL._
             (for {
@@ -340,7 +340,7 @@ class AlertSrv(
             Future.successful(())
         }
       }
-        .map(_ => caze)
+      .map(_ ⇒ caze)
     updatedCase.onComplete { _ ⇒
       // remove temporary files
       artifactsFields