From d200e3c4aa1dfcf66d8521db99870414cdc6f620 Mon Sep 17 00:00:00 2001
From: To-om <toom@thehive-project.org>
Date: Wed, 13 Mar 2019 15:34:48 +0100
Subject: [PATCH] #175 enrich worker info from definition during migration

---
 app/org/thp/cortex/models/Migration.scala   | 29 ++++++++++++++++++---
 app/org/thp/cortex/services/WorkerSrv.scala |  8 +++---
 2 files changed, 29 insertions(+), 8 deletions(-)

diff --git a/app/org/thp/cortex/models/Migration.scala b/app/org/thp/cortex/models/Migration.scala
index ac363dc61..62ce49c88 100644
--- a/app/org/thp/cortex/models/Migration.scala
+++ b/app/org/thp/cortex/models/Migration.scala
@@ -4,9 +4,10 @@ import javax.inject.{ Inject, Singleton }
 import scala.concurrent.{ ExecutionContext, Future }
 import scala.util.Success
 
-import play.api.libs.json.{ JsString, Json }
+import play.api.Logger
+import play.api.libs.json.{ JsNull, JsString, JsValue, Json }
 
-import org.thp.cortex.services.{ OrganizationSrv, UserSrv }
+import org.thp.cortex.services.{ OrganizationSrv, UserSrv, WorkerSrv }
 
 import org.elastic4play.controllers.Fields
 import org.elastic4play.services.Operation._
@@ -17,8 +18,10 @@ import org.elastic4play.utils.Hasher
 class Migration @Inject() (
     userSrv: UserSrv,
     organizationSrv: OrganizationSrv,
+    workerSrv: WorkerSrv,
     implicit val ec: ExecutionContext) extends MigrationOperations {
 
+  lazy val logger = Logger(getClass)
   def beginMigration(version: Int): Future[Unit] = Future.successful(())
 
   def endMigration(version: Int): Future[Unit] = {
@@ -27,7 +30,7 @@ class Migration @Inject() (
         "name" → "cortex",
         "description" → "Default organization",
         "status" → "Active")))
-        .transform { case _ ⇒ Success(()) } // ignore errors (already exist)
+        .transform(_ ⇒ Success(())) // ignore errors (already exist)
     }
   }
 
@@ -61,6 +64,24 @@ class Migration @Inject() (
         renameEntity("analyzerConfig", "workerConfig"),
         addAttribute("workerConfig", "type" → JsString(WorkerType.analyzer.toString)))
 
-    case _ ⇒ Nil
+    case DatabaseState(2) ⇒
+      Seq(mapEntity("worker") { worker ⇒
+        val definitionId = (worker \ "workerDefinitionId").asOpt[String]
+        definitionId
+          .flatMap(workerSrv.getDefinition(_).toOption)
+          .fold {
+            logger.warn(s"no definition found for worker ${definitionId.getOrElse(worker)}. You should probably have to disable and re-enable it")
+            worker
+          } { definition ⇒
+            worker +
+              ("version" -> JsString(definition.version)) +
+              ("author" -> JsString(definition.author)) +
+              ("url" -> JsString(definition.url)) +
+              ("license" -> JsString(definition.license)) +
+              ("command" -> definition.command.fold[JsValue](JsNull)(c ⇒ JsString(c.toString))) +
+              ("dockerImage" -> definition.image.fold[JsValue](JsNull)(JsString.apply)) +
+              ("baseConfig" -> definition.baseConfiguration.fold[JsValue](JsNull)(JsString.apply))
+          }
+      })
   }
 }
diff --git a/app/org/thp/cortex/services/WorkerSrv.scala b/app/org/thp/cortex/services/WorkerSrv.scala
index 9207770cf..1a9c2c584 100644
--- a/app/org/thp/cortex/services/WorkerSrv.scala
+++ b/app/org/thp/cortex/services/WorkerSrv.scala
@@ -6,7 +6,7 @@ import java.nio.file.{ Files, Path, Paths }
 import javax.inject.{ Inject, Singleton }
 import scala.collection.JavaConverters._
 import scala.concurrent.{ ExecutionContext, Future }
-import scala.util.Try
+import scala.util.{ Failure, Success, Try }
 
 import play.api.libs.json.{ JsObject, JsString, Json }
 import play.api.libs.ws.WSClient
@@ -53,9 +53,9 @@ class WorkerSrv @Inject() (
 
   rescan()
 
-  def getDefinition(workerId: String): Future[WorkerDefinition] = workerMap.get(workerId) match {
-    case Some(worker) ⇒ Future.successful(worker)
-    case None         ⇒ Future.failed(NotFoundError(s"Worker $workerId not found"))
+  def getDefinition(workerId: String): Try[WorkerDefinition] = workerMap.get(workerId) match {
+    case Some(worker) ⇒ Success(worker)
+    case None         ⇒ Failure(NotFoundError(s"Worker $workerId not found"))
   }
 
   //  def listDefinitions: (Source[WorkerDefinition, NotUsed], Future[Long]) = Source(workerMap.values.toList) → Future.successful(workerMap.size.toLong)