#869 Use observable TLP if possible when executing a responder

TheHive-Project · Feb 13, 2019 · e179b6e · e179b6e
1 parent 4411840
commit e179b6e
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 4 deletions.
diff --git a/thehive-cortex/app/connectors/cortex/services/ActionOperation.scala b/thehive-cortex/app/connectors/cortex/services/ActionOperation.scala
@@ -139,6 +139,18 @@ class ActionOperationSrv @Inject() (
     }
   }
 
+  def findArtifactEntity(entity: BaseEntity): Future[Artifact] = {
+    import org.elastic4play.services.QueryDSL._
+
+    (entity, entity.model) match {
+      case (a: Artifact, _) ⇒ Future.successful(a)
+      case (_, model: ChildModelDef[_, _, _, _]) ⇒
+        findSrv(model.parentModel, "_id" ~= entity.parentId.getOrElse(throw InternalError(s"Child entity $entity has no parent ID")), Some("0-1"), Nil)
+          ._1.runWith(Sink.head).flatMap(findArtifactEntity _)
+      case _ ⇒ Future.failed(BadRequestError("Artifact not found"))
+    }
+  }
+
   def execute(entity: BaseEntity, operation: ActionOperation)(implicit authContext: AuthContext): Future[ActionOperation] = {
     if (operation.status == ActionOperationStatus.Waiting) {
       Retry()(classOf[VersionConflictEngineException]) {

diff --git a/thehive-cortex/app/connectors/cortex/services/CortexActionSrv.scala b/thehive-cortex/app/connectors/cortex/services/CortexActionSrv.scala
@@ -72,10 +72,14 @@ class CortexActionSrv @Inject() (
 
   def findResponderFor(entityType: String, entityId: String): Future[Seq[Responder]] = {
     for {
-      (tlp, pap) ← getEntity(entityType, entityId)
-        .flatMap(actionOperationSrv.findCaseEntity)
-        .map { caze ⇒ (caze.tlp(), caze.pap()) }
-        .recover { case _ ⇒ (0L, 0L) }
+      entity ← getEntity(entityType, entityId)
+      artifactTlp ← actionOperationSrv
+        .findArtifactEntity(entity)
+        .map(a ⇒ Some(a.tlp()))
+        .recover { case _ ⇒ None }
+      (tlp, pap) ← actionOperationSrv.findCaseEntity(entity)
+        .map { caze ⇒ (artifactTlp.getOrElse(caze.tlp()), caze.pap()) }
+        .recover { case _ ⇒ (artifactTlp.getOrElse(0L), 0L) }
       query = Json.obj(
         "dataTypeList" → s"thehive:$entityType")
       responders ← findResponders(query)