#616 Add PAP attribute in case

TheHive-Project · Jul 6, 2018 · 260b200 · 260b200
1 parent 6faceec
commit 260b200
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 14 deletions.
diff --git a/thehive-backend/app/models/Case.scala b/thehive-backend/app/models/Case.scala
@@ -44,6 +44,7 @@ trait CaseAttributes { _: AttributeDef ⇒
   val tags: A[Seq[String]] = multiAttribute("tags", F.stringFmt, "Case tags")
   val flag: A[Boolean] = attribute("flag", F.booleanFmt, "Flag of the case", false)
   val tlp: A[Long] = attribute("tlp", TlpAttributeFormat, "TLP level", 2L)
+  val pap: A[Long] = attribute("pap", TlpAttributeFormat, "PAP level", 2L)
   val status: A[CaseStatus.Value] = attribute("status", F.enumFmt(CaseStatus), "Status of the case", CaseStatus.Open)
   val metrics: A[JsValue] = attribute("metrics", F.metricsFmt, "List of metrics", JsObject.empty)
   val resolutionStatus: A[Option[CaseResolutionStatus.Value]] = optionalAttribute("resolutionStatus", F.enumFmt(CaseResolutionStatus), "Resolution status of the case")

diff --git a/thehive-cortex/app/connectors/cortex/controllers/CortexCtrl.scala b/thehive-cortex/app/connectors/cortex/controllers/CortexCtrl.scala
@@ -167,9 +167,7 @@ class CortexCtrl @Inject() (
   }
 
   def getWorkers(entityType: String, entityId: String): Action[AnyContent] = authenticated(Roles.read).async { implicit request ⇒
-    val query = Json.obj(
-      "dataTypeList" -> s"thehive:$entityType")
-    cortexActionSrv.findWorkers(query).map { workers ⇒
+    cortexActionSrv.findWorkerFor(entityType, entityId).map { workers ⇒
       renderer.toOutput(OK, workers)
     }
   }

diff --git a/thehive-cortex/app/connectors/cortex/services/CortexActionSrv.scala b/thehive-cortex/app/connectors/cortex/services/CortexActionSrv.scala
@@ -66,6 +66,19 @@ class CortexActionSrv @Inject() (
       }
   }
 
+  def findWorkerFor(entityType: String, entityId: String): Future[Seq[Worker]] = {
+    for {
+      (tlp, pap) ← getEntity(entityType, entityId)
+        .flatMap(actionOperationSrv.findCaseEntity)
+        .map { caze ⇒ (caze.tlp(), caze.pap()) }
+        .recover { case _ ⇒ (0L, 0L) }
+      query = Json.obj(
+        "dataTypeList" -> s"thehive:$entityType")
+      workers ← findWorkers(query)
+      applicableWorkers = workers.filter(w ⇒ w.maxTlp.fold(true)(_ >= tlp) && w.maxPap.fold(true)(_ >= pap))
+    } yield applicableWorkers
+  }
+
   def find(queryDef: QueryDef, range: Option[String], sortBy: Seq[String]): (Source[Action, NotUsed], Future[Long]) = {
     findSrv[ActionModel, Action](actionModel, queryDef, range, sortBy)
   }
@@ -125,6 +138,17 @@ class CortexActionSrv @Inject() (
     ()
   }
 
+  def getEntity(objectType: String, objectId: String): Future[BaseEntity] = {
+    import org.elastic4play.services.QueryDSL._
+
+    findSrv.apply(Some(objectType), withId(objectId), Some("0-1"), Nil)._1
+      .runWith(Sink.headOption)
+      .flatMap {
+        case Some(entity) ⇒ Future.successful(entity)
+        case None         ⇒ Future.failed(NotFoundError(s"$objectType $objectId not found"))
+      }
+  }
+
   def executeAction(fields: Fields)(implicit authContext: AuthContext): Future[Action] = {
     def getWorker(cortexClient: CortexClient): Future[Worker] = {
       fields.getString("workerId").map(cortexClient.getWorkerById) orElse
@@ -149,17 +173,6 @@ class CortexActionSrv @Inject() (
         }
     }
 
-    def getEntity(objectType: String, objectId: String): Future[BaseEntity] = {
-      import org.elastic4play.services.QueryDSL._
-
-      findSrv.apply(Some(objectType), withId(objectId), Some("0-1"), Nil)._1
-        .runWith(Sink.headOption)
-        .flatMap {
-          case Some(entity) ⇒ Future.successful(entity)
-          case None         ⇒ Future.failed(NotFoundError(s"$objectType $objectId not found"))
-        }
-    }
-
     for {
       objectType ← fields.getString("objectType").fold[Future[String]](Future.failed(MissingAttributeError("action.objectType")))(Future.successful)
       objectId ← fields.getString("objectId").fold[Future[String]](Future.failed(MissingAttributeError("action.objectId")))(Future.successful)