From fc938f7f9f668d8ec8660f491cd18799bd733339 Mon Sep 17 00:00:00 2001
From: To-om <toom@thehive-project.org>
Date: Thu, 26 Jul 2018 09:38:17 +0200
Subject: [PATCH] #110 Include type in worker id computation in order to
 prevent collision between responders and analyzers

---
 app/org/thp/cortex/models/Migration.scala | 40 +++++++++++++++--------
 app/org/thp/cortex/models/Worker.scala    |  3 +-
 2 files changed, 28 insertions(+), 15 deletions(-)

diff --git a/app/org/thp/cortex/models/Migration.scala b/app/org/thp/cortex/models/Migration.scala
index bdb9fca0a..503f8b627 100644
--- a/app/org/thp/cortex/models/Migration.scala
+++ b/app/org/thp/cortex/models/Migration.scala
@@ -11,6 +11,7 @@ import org.thp.cortex.services.{ OrganizationSrv, UserSrv }
 import org.elastic4play.controllers.Fields
 import org.elastic4play.services.Operation._
 import org.elastic4play.services.{ DatabaseState, MigrationOperations, Operation }
+import org.elastic4play.utils.Hasher
 
 @Singleton
 class Migration @Inject() (
@@ -31,20 +32,31 @@ class Migration @Inject() (
   }
 
   val operations: PartialFunction[DatabaseState, Seq[Operation]] = {
-    case DatabaseState(1) ⇒ Seq(
-      // add type to analyzer
-      addAttribute("analyzer", "type" → JsString("analyzer")),
-
-      renameAttribute("job", "workerDefinitionId", "analyzerDefinitionId"),
-      renameAttribute("job", "workerId", "analyzerId"),
-      renameAttribute("job", "workerName", "analyzerName"),
-
-      renameEntity("analyzer", "worker"),
-      renameAttribute("worker", "workerDefinitionId", "analyzerDefinitionId"),
-      addAttribute("worker", "type" → JsString(WorkerType.analyzer.toString)),
-
-      renameEntity("analyzerConfig", "workerConfig"),
-      addAttribute("workerConfig", "type" → JsString(WorkerType.analyzer.toString)))
+    case DatabaseState(1) ⇒
+      val hasher = Hasher("MD5")
+      Seq(
+
+        // add type to analyzer
+        addAttribute("analyzer", "type" → JsString("analyzer")),
+
+        renameAttribute("job", "workerDefinitionId", "analyzerDefinitionId"),
+        renameAttribute("job", "workerId", "analyzerId"),
+        renameAttribute("job", "workerName", "analyzerName"),
+
+        renameEntity("analyzer", "worker"),
+        renameAttribute("worker", "workerDefinitionId", "analyzerDefinitionId"),
+        addAttribute("worker", "type" → JsString(WorkerType.analyzer.toString)),
+        mapEntity("worker") { worker ⇒
+          val id = for {
+            organizationId ← (worker \ "_parent").asOpt[String]
+            name ← (worker \ "name").asOpt[String]
+            tpe ← (worker \ "type").asOpt[String]
+          } yield hasher.fromString(s"${organizationId}_${name}_$tpe").head.toString
+          worker + ("_id" → JsString(id.getOrElse("<null>")))
+        },
+
+        renameEntity("analyzerConfig", "workerConfig"),
+        addAttribute("workerConfig", "type" → JsString(WorkerType.analyzer.toString)))
 
     case _ ⇒ Nil
   }
diff --git a/app/org/thp/cortex/models/Worker.scala b/app/org/thp/cortex/models/Worker.scala
index 0dcac0d01..5caf6d679 100644
--- a/app/org/thp/cortex/models/Worker.scala
+++ b/app/org/thp/cortex/models/Worker.scala
@@ -45,7 +45,8 @@ class WorkerModel @Inject() (organizationModel: OrganizationModel) extends Child
     val id = for {
       organizationId ← parent.map(_.id)
       name ← (attrs \ "name").asOpt[String]
-    } yield hasher.fromString(s"${organizationId}_$name").head.toString
+      tpe ← (attrs \ "type").asOpt[String]
+    } yield hasher.fromString(s"${organizationId}_${name}_$tpe").head.toString
     Future.successful(attrs + ("_id" → JsString(id.getOrElse("<null>"))))
   }
 }