#175 Don't use definition to render worker output

TheHive-Project · Mar 13, 2019 · a719f5b · a719f5b
1 parent ed207e6
commit a719f5b
Show file tree

Hide file tree

Showing 4 changed files with 38 additions and 96 deletions.
diff --git a/app/org/thp/cortex/controllers/AnalyzerCtrl.scala b/app/org/thp/cortex/controllers/AnalyzerCtrl.scala
@@ -1,18 +1,16 @@
 package org.thp.cortex.controllers
 
-import scala.concurrent.ExecutionContext
+import scala.concurrent.{ ExecutionContext, Future }
 
-import play.api.libs.json.{ JsNumber, JsObject, JsString, Json }
+import play.api.libs.json.{ JsObject, JsString, Json }
 import play.api.mvc.{ AbstractController, Action, AnyContent, ControllerComponents }
 
 import akka.stream.Materializer
-import akka.stream.scaladsl.Sink
 import javax.inject.{ Inject, Singleton }
-import org.thp.cortex.models.{ Roles, Worker, WorkerDefinition }
+import org.thp.cortex.models.{ Roles, Worker }
 import org.thp.cortex.services.{ UserSrv, WorkerSrv }
 
 import org.elastic4play.controllers.{ Authenticated, Fields, FieldsBodyParser, Renderer }
-import org.elastic4play.models.JsonFormat.baseModelEntityWrites
 import org.elastic4play.services.JsonFormat.queryReads
 import org.elastic4play.services.{ QueryDSL, QueryDef }
 
@@ -42,28 +40,6 @@ class AnalyzerCtrl @Inject() (
       .map(a ⇒ renderer.toOutput(OK, analyzerJson(isAdmin)(a)))
   }
 
-  private val emptyAnalyzerDefinitionJson = Json.obj(
-    "version" → "0.0",
-    "description" → "unknown",
-    "dataTypeList" → Nil,
-    "author" → "unknown",
-    "url" → "unknown",
-    "license" → "unknown")
-
-  private def analyzerJson(analyzer: Worker, analyzerDefinition: Option[WorkerDefinition]) = {
-    analyzer.toJson ++ analyzerDefinition.fold(emptyAnalyzerDefinitionJson) { ad ⇒
-      Json.obj(
-        "maxTlp" → (analyzer.config \ "max_tlp").asOpt[JsNumber],
-        "maxPap" → (analyzer.config \ "max_pap").asOpt[JsNumber],
-        "version" → ad.version,
-        "description" → ad.description,
-        "author" → ad.author,
-        "url" → ad.url,
-        "license" → ad.license,
-        "baseConfig" → ad.baseConfiguration)
-    } + ("analyzerDefinitionId" → JsString(analyzer.workerDefinitionId())) // For compatibility reason
-  }
-
   private def analyzerJson(isAdmin: Boolean)(analyzer: Worker): JsObject = {
     if (isAdmin)
       analyzer.toJson + ("configuration" → Json.parse(analyzer.configuration())) + ("analyzerDefinitionId" → JsString(analyzer.workerDefinitionId()))
@@ -73,30 +49,24 @@ class AnalyzerCtrl @Inject() (
 
   def listForType(dataType: String): Action[AnyContent] = authenticated(Roles.read).async { request ⇒
     import org.elastic4play.services.QueryDSL._
-    workerSrv.findAnalyzersForUser(request.userId, "dataTypeList" ~= dataType, Some("all"), Nil)
-      ._1
-      .mapAsyncUnordered(2) { analyzer ⇒
-        workerSrv.getDefinition(analyzer.workerDefinitionId())
-          .map(ad ⇒ analyzerJson(analyzer, Some(ad)))
-      }
-      .runWith(Sink.seq)
-      .map(analyzers ⇒ renderer.toOutput(OK, analyzers))
+    val (responderList, responderCount) = workerSrv.findAnalyzersForUser(request.userId, "dataTypeList" ~= dataType, Some("all"), Nil)
+    renderer.toOutput(OK, responderList.map(analyzerJson(isAdmin = false)), responderCount)
   }
 
   def create(analyzerDefinitionId: String): Action[Fields] = authenticated(Roles.orgAdmin).async(fieldsBodyParser) { implicit request ⇒
     for {
       organizationId ← userSrv.getOrganizationId(request.userId)
-      workerDefinition ← workerSrv.getDefinition(analyzerDefinitionId)
+      workerDefinition ← Future.fromTry(workerSrv.getDefinition(analyzerDefinitionId))
       analyzer ← workerSrv.create(organizationId, workerDefinition, request.body)
-    } yield renderer.toOutput(CREATED, analyzerJson(analyzer, Some(workerDefinition)))
+    } yield renderer.toOutput(CREATED, analyzerJson(isAdmin = false)(analyzer))
   }
 
-  def listDefinitions: Action[AnyContent] = authenticated(Roles.orgAdmin, Roles.superAdmin).async { implicit request ⇒
+  def listDefinitions: Action[AnyContent] = authenticated(Roles.orgAdmin, Roles.superAdmin).async { _ ⇒
     val (analyzers, analyzerTotal) = workerSrv.listAnalyzerDefinitions
     renderer.toOutput(OK, analyzers, analyzerTotal)
   }
 
-  def scan: Action[AnyContent] = authenticated(Roles.orgAdmin, Roles.superAdmin) { implicit request ⇒
+  def scan: Action[AnyContent] = authenticated(Roles.orgAdmin, Roles.superAdmin) { _ ⇒
     workerSrv.rescan()
     NoContent
   }

diff --git a/app/org/thp/cortex/controllers/ResponderCtrl.scala b/app/org/thp/cortex/controllers/ResponderCtrl.scala
@@ -6,7 +6,6 @@ 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 javax.inject.{ Inject, Singleton }
 import org.thp.cortex.models.{ Roles, Worker, WorkerDefinition }
 import org.thp.cortex.services.{ UserSrv, WorkerSrv }
@@ -33,15 +32,13 @@ class ResponderCtrl @Inject() (
     val sort = request.body.getStrings("sort").getOrElse(Nil)
     val isAdmin = request.roles.contains(Roles.orgAdmin)
     val (responders, responderTotal) = workerSrv.findRespondersForUser(request.userId, query, range, sort)
-    val enrichedResponders = responders.mapAsync(2)(responderJson(isAdmin))
-    renderer.toOutput(OK, enrichedResponders, responderTotal)
+    renderer.toOutput(OK, responders.map(responderJson(isAdmin)), responderTotal)
   }
 
   def get(responderId: String): Action[AnyContent] = authenticated(Roles.read).async { request ⇒
     val isAdmin = request.roles.contains(Roles.orgAdmin)
     workerSrv.getForUser(request.userId, responderId)
-      .flatMap(responderJson(isAdmin))
-      .map(renderer.toOutput(OK, _))
+      .map(responder ⇒ renderer.toOutput(OK, responderJson(isAdmin)(responder)))
   }
 
   private val emptyResponderDefinitionJson = Json.obj(
@@ -66,32 +63,23 @@ class ResponderCtrl @Inject() (
     }
   }
 
-  private def responderJson(isAdmin: Boolean)(responder: Worker): Future[JsObject] = {
-    workerSrv.getDefinition(responder.workerDefinitionId())
-      .map(responderDefinition ⇒ responderJson(responder, Some(responderDefinition)))
-      .recover { case _ ⇒ responderJson(responder, None) }
-      .map {
-        case a if isAdmin ⇒ a + ("configuration" → Json.parse(responder.configuration()))
-        case a            ⇒ a
-      }
+  private def responderJson(isAdmin: Boolean)(responder: Worker): JsObject = {
+    if (isAdmin)
+      responder.toJson + ("configuration" → Json.parse(responder.configuration()))
+    else
+      responder.toJson
   }
 
   def listForType(dataType: String): Action[AnyContent] = authenticated(Roles.read).async { request ⇒
     import org.elastic4play.services.QueryDSL._
-    workerSrv.findRespondersForUser(request.userId, "dataTypeList" ~= dataType, Some("all"), Nil)
-      ._1
-      .mapAsyncUnordered(2) { responder ⇒
-        workerSrv.getDefinition(responder.workerDefinitionId())
-          .map(ad ⇒ responderJson(responder, Some(ad)))
-      }
-      .runWith(Sink.seq)
-      .map(responders ⇒ renderer.toOutput(OK, responders))
+    val (responderList, responderCount) = workerSrv.findRespondersForUser(request.userId, "dataTypeList" ~= dataType, Some("all"), Nil)
+    renderer.toOutput(OK, responderList.map(responderJson(false)), responderCount)
   }
 
   def create(responderDefinitionId: String): Action[Fields] = authenticated(Roles.orgAdmin).async(fieldsBodyParser) { implicit request ⇒
     for {
       organizationId ← userSrv.getOrganizationId(request.userId)
-      workerDefinition ← workerSrv.getDefinition(responderDefinitionId)
+      workerDefinition ← Future.fromTry(workerSrv.getDefinition(responderDefinitionId))
       responder ← workerSrv.create(organizationId, workerDefinition, request.body)
     } yield renderer.toOutput(CREATED, responderJson(responder, Some(workerDefinition)))
   }
@@ -117,7 +105,6 @@ class ResponderCtrl @Inject() (
     for {
       responder ← workerSrv.getForUser(request.userId, responderId)
       updatedResponder ← workerSrv.update(responder, request.body)
-      updatedResponderJson ← responderJson(isAdmin = true)(updatedResponder)
-    } yield renderer.toOutput(OK, updatedResponderJson)
+    } yield renderer.toOutput(OK, responderJson(isAdmin = true)(updatedResponder))
   }
 }
diff --git a/app/org/thp/cortex/services/JobRunnerSrv.scala b/app/org/thp/cortex/services/JobRunnerSrv.scala
@@ -175,22 +175,13 @@ class JobRunnerSrv @Inject() (
     }
   }
 
-  //  private def fixArtifact(artifact: Fields): Fields = {
-  //    def rename(oldName: String, newName: String): Fields ⇒ Fields = fields ⇒
-  //      fields.getValue(oldName).fold(fields)(v ⇒ fields.unset(oldName).set(newName, v))
-  //
-  //    rename("value", "data").andThen(
-  //      rename("type", "dataType"))(artifact)
-  //  }
-
   def run(worker: Worker, job: Job)(implicit authContext: AuthContext): Future[Job] = {
     prepareJobFolder(worker, job).flatMap { jobFolder ⇒
+      val executionContext = worker.tpe match {
+        case WorkerType.analyzer  ⇒ analyzerExecutionContext
+        case WorkerType.responder ⇒ responderExecutionContext
+      }
       val finishedJob = for {
-        workerDefinition ← workerSrv.getDefinition(worker.workerDefinitionId())
-        executionContext = workerDefinition.tpe match {
-          case WorkerType.analyzer  ⇒ analyzerExecutionContext
-          case WorkerType.responder ⇒ responderExecutionContext
-        }
         _ ← startJob(job)
         j ← runners
           .foldLeft[Option[Future[Unit]]](None) {
@@ -223,7 +214,7 @@ class JobRunnerSrv @Inject() (
             error ⇒ endJob(job, JobStatus.Failure, Option(error.getMessage), Some(readFile(jobFolder.resolve("input").resolve("input.json")))),
             _ ⇒ extractReport(jobFolder, job))
         }
-      //.andThen { case _ ⇒ delete(jobFolder) }
+        .andThen { case _ ⇒ delete(jobFolder) }
     }
   }
 

diff --git a/app/org/thp/cortex/services/MispSrv.scala b/app/org/thp/cortex/services/MispSrv.scala
@@ -35,25 +35,19 @@ class MispSrv @Inject() (
     val (analyzers, analyzerCount) = workerSrv.findAnalyzersForUser(authContext.userId, QueryDSL.any, Some("all"), Nil)
 
     val mispAnalyzers = analyzers
-      .mapAsyncUnordered(1) { analyzer ⇒
-        workerSrv.getDefinition(analyzer.workerDefinitionId())
-          .map(ad ⇒ Some(analyzer → ad))
-          .recover { case _ ⇒ None }
-      }
-      .collect {
-        case Some((analyzer, analyzerDefinition)) ⇒
-          Json.obj(
-            "name" → analyzer.name(),
-            "type" → "cortex",
-            "mispattributes" → Json.obj(
-              "input" → analyzer.dataTypeList().flatMap(dataType2mispType).distinct,
-              "output" → Json.arr()),
-            "meta" → Json.obj(
-              "module-type" → Json.arr("cortex"),
-              "description" → analyzer.description(),
-              "author" → analyzerDefinition.author,
-              "version" → analyzerDefinition.version,
-              "config" → Json.arr()))
+      .map { analyzer ⇒
+        Json.obj(
+          "name" → analyzer.name(),
+          "type" → "cortex",
+          "mispattributes" → Json.obj(
+            "input" → analyzer.dataTypeList().flatMap(dataType2mispType).distinct,
+            "output" → Json.arr()),
+          "meta" → Json.obj(
+            "module-type" → Json.arr("cortex"),
+            "description" → analyzer.description(),
+            "author" → analyzer.author(),
+            "version" → analyzer.vers(),
+            "config" → Json.arr()))
       }
     mispAnalyzers → analyzerCount
   }