From e287c51a5acfb30cadfbcb19f910e7898414cf7d Mon Sep 17 00:00:00 2001
From: To-om <toom@thehive-project.org>
Date: Tue, 14 Jan 2020 16:10:01 +0100
Subject: [PATCH] #234 Invalid disapeared analyzers/responders

---
 app/org/thp/cortex/services/WorkerSrv.scala | 35 ++++++++++++---------
 1 file changed, 21 insertions(+), 14 deletions(-)

diff --git a/app/org/thp/cortex/services/WorkerSrv.scala b/app/org/thp/cortex/services/WorkerSrv.scala
index cb8203a18..20040592d 100644
--- a/app/org/thp/cortex/services/WorkerSrv.scala
+++ b/app/org/thp/cortex/services/WorkerSrv.scala
@@ -3,27 +3,24 @@ package org.thp.cortex.services
 import java.net.URL
 import java.nio.file.{Files, Path, Paths}
 
-import javax.inject.{Inject, Provider, Singleton}
-import scala.collection.JavaConverters._
-import scala.concurrent.{ExecutionContext, Future}
-import scala.util.{Failure, Success, Try}
-
-import play.api.libs.json.{JsObject, JsString, Json}
-import play.api.libs.ws.WSClient
-import play.api.{Configuration, Logger}
-
 import akka.NotUsed
 import akka.stream.Materializer
 import akka.stream.scaladsl.{Sink, Source}
-import org.thp.cortex.models._
-
+import javax.inject.{Inject, Provider, Singleton}
 import org.elastic4play._
 import org.elastic4play.controllers.{Fields, StringInputValue}
+import org.elastic4play.database.ModifyConfig
 import org.elastic4play.services._
-import org.scalactic._
 import org.scalactic.Accumulation._
+import org.scalactic._
+import org.thp.cortex.models._
+import play.api.libs.json.{JsArray, JsObject, JsString, Json}
+import play.api.libs.ws.WSClient
+import play.api.{Configuration, Logger}
 
-import org.elastic4play.database.ModifyConfig
+import scala.collection.JavaConverters._
+import scala.concurrent.{ExecutionContext, Future}
+import scala.util.{Failure, Success, Try}
 
 @Singleton
 class WorkerSrv @Inject()(
@@ -128,11 +125,21 @@ class WorkerSrv @Inject()(
   private def find(queryDef: QueryDef, range: Option[String], sortBy: Seq[String]): (Source[Worker, NotUsed], Future[Long]) =
     findSrv[WorkerModel, Worker](workerModel, queryDef, range, sortBy)
 
-  def rescan(): Unit =
+  def rescan(): Unit = {
+    import org.elastic4play.services.QueryDSL._
     scan(
       analyzersURLs.map(_    → WorkerType.analyzer) ++
         respondersURLs.map(_ → WorkerType.responder)
     )
+    userSrv.inInitAuthContext { implicit authContext ⇒
+      find(any, Some("all"), Nil)._1.runForeach { worker ⇒
+        workerMap.get(worker.id) match {
+          case Some(wd) ⇒ update(worker, Fields.empty.set("dataTypeList", Json.toJson(wd.dataTypeList)))
+          case None     ⇒ update(worker, Fields.empty.set("dataTypeList", JsArray.empty))
+        }
+      }
+    }
+  }
 
   def scan(workerUrls: Seq[(String, WorkerType.Type)]): Unit = {
     def readUrl(url: URL, workerType: WorkerType.Type): Future[Seq[WorkerDefinition]] =