#53 Fix Cortex id when several Cortex is configured

TheHive-Project · Dec 12, 2016 · b6c3f8b · b6c3f8b
1 parent 78c0687
commit b6c3f8b
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 18 deletions.
diff --git a/thehive-cortex/app/connectors/cortex/models/Analyzer.scala b/thehive-cortex/app/connectors/cortex/models/Analyzer.scala
@@ -1,15 +1,10 @@
 package connectors.cortex.models
 
-trait CortexModel[O] { self ⇒
-  def onCortex(cortexId: String): O
-}
-
 case class Analyzer(
     name: String,
     version: String,
     description: String,
     dataTypeList: Seq[String],
-    cortexIds: List[String] = Nil) extends CortexModel[Analyzer] {
+    cortexIds: List[String] = Nil) {
   def id = (name + "_" + version).replaceAll("\\.", "_")
-  def onCortex(cortexId: String) = copy(cortexIds = cortexId :: cortexIds)
 }
diff --git a/thehive-cortex/app/connectors/cortex/models/Job.scala b/thehive-cortex/app/connectors/cortex/models/Job.scala
@@ -44,6 +44,6 @@ class Job(model: JobModel, attributes: JsObject) extends EntityDef[JobModel, Job
   override def toJson = super.toJson + ("report" → report().fold[JsValue](JsObject(Nil))(r ⇒ Json.parse(r))) // FIXME is parse fails (invalid report)
 }
 
-case class CortexJob(id: String, analyzerId: String, artifact: CortexArtifact, date: Date, status: JobStatus.Type, cortexIds: List[String] = Nil) extends CortexModel[CortexJob] {
+case class CortexJob(id: String, analyzerId: String, artifact: CortexArtifact, date: Date, status: JobStatus.Type, cortexIds: List[String] = Nil) {
   def onCortex(cortexId: String) = copy(cortexIds = cortexId :: cortexIds)
 }
diff --git a/thehive-cortex/app/connectors/cortex/services/CortexSrv.scala b/thehive-cortex/app/connectors/cortex/services/CortexSrv.scala
@@ -25,7 +25,7 @@ import org.elastic4play.{ InternalError, NotFoundError }
 import org.elastic4play.controllers.Fields
 import org.elastic4play.services.{ AttachmentSrv, AuthContext, CreateSrv, EventSrv, FindSrv, GetSrv, QueryDef, UpdateSrv }
 
-import connectors.cortex.models.{ Analyzer, CortexJob, CortexModel, DataArtifact, FileArtifact, Job, JobModel, JobStatus }
+import connectors.cortex.models.{ Analyzer, CortexJob, DataArtifact, FileArtifact, Job, JobModel, JobStatus }
 import connectors.cortex.models.JsonFormat.{ cortexJobFormat, jobStatusFormat }
 import models.Artifact
 import services.{ ArtifactSrv, MergeArtifact }
@@ -104,14 +104,6 @@ class CortexSrv @Inject() (
     findSrv[JobModel, Job](jobModel, queryDef, range, sortBy)
   }
 
-  def askAllCortex[A](f: CortexClient ⇒ Future[Seq[CortexModel[A]]]): Future[Seq[A]] = {
-    Future
-      .traverse(cortexConfig.instances.toSeq) {
-        case cortex ⇒ f(cortex).map(_.map(_.onCortex(cortex.name)))
-      }
-      .map(_.flatten)
-  }
-
   def getAnalyzer(analyzerId: String): Future[Analyzer] = {
     Future
       .traverse(cortexConfig.instances) {
@@ -126,12 +118,42 @@ class CortexSrv @Inject() (
       }
   }
 
+  def askAnalyzersOnAllCortex(f: CortexClient ⇒ Future[Seq[Analyzer]]): Future[Seq[Analyzer]] = {
+    Future
+      .traverse(cortexConfig.instances) {
+        case cortex ⇒ f(cortex)
+      }
+      .map(_.flatten)
+  }
+
   def getAnalyzersFor(dataType: String): Future[Seq[Analyzer]] = {
-    askAllCortex(_.listAnalyzerForType(dataType))
+    Future
+      .traverse(cortexConfig.instances) {
+        case cortex ⇒ cortex.listAnalyzerForType(dataType)
+      }
+      .map { listOfListOfAnalyzers ⇒
+        val analysers = listOfListOfAnalyzers.flatten
+        analysers
+          .groupBy(_.name)
+          .values
+          .map(_.reduce((a1, a2) ⇒ a1.copy(cortexIds = a1.cortexIds ::: a2.cortexIds)))
+          .toSeq
+      }
   }
 
   def listAnalyzer: Future[Seq[Analyzer]] = {
-    askAllCortex(_.listAnalyzer)
+    Future
+      .traverse(cortexConfig.instances) {
+        case cortex ⇒ cortex.listAnalyzer
+      }
+      .map { listOfListOfAnalyzers ⇒
+        val analysers = listOfListOfAnalyzers.flatten
+        analysers
+          .groupBy(_.name)
+          .values
+          .map(_.reduce((a1, a2) ⇒ a1.copy(cortexIds = a1.cortexIds ::: a2.cortexIds)))
+          .toSeq
+      }
   }
 
   def getJob(jobId: String): Future[Job] = {