#1340 Improve migration

TheHive-Project · Jun 8, 2020 · f1203b1 · f1203b1
1 parent 6c9c03a
commit f1203b1
Show file tree

Hide file tree

Showing 6 changed files with 662 additions and 389 deletions.
diff --git a/conf/migration-logback.xml b/conf/migration-logback.xml
@@ -46,7 +46,6 @@
     -->
     <logger name="org.janusgraph.graphdb.transaction.StandardJanusGraphTx" level="ERROR"/>
     <logger name="org.thp.thehive" level="INFO"/>
-    <logger name="org.elastic4play.services" level="DEBUG"/>
 
     <root level="INFO">
         <appender-ref ref="ASYNCFILE"/>

diff --git a/migration/src/main/scala/org/thp/thehive/migration/Input.scala b/migration/src/main/scala/org/thp/thehive/migration/Input.scala
@@ -5,6 +5,9 @@ import akka.stream.scaladsl.Source
 import com.typesafe.config.Config
 import org.thp.thehive.migration.dto._
 
+import scala.concurrent.Future
+import scala.util.Try
+
 case class Filter(caseFromDate: Long, alertFromDate: Long, auditFromDate: Long)
 
 object Filter {
@@ -21,32 +24,60 @@ object Filter {
 }
 
 trait Input {
-  def listOrganisations(filter: Filter): Source[InputOrganisation, NotUsed]
-  def listCases(filter: Filter): Source[InputCase, NotUsed]
-  def listCaseObservables(filter: Filter): Source[(String, InputObservable), NotUsed]
-  def listCaseObservables(caseId: String): Source[(String, InputObservable), NotUsed]
-  def listCaseTasks(filter: Filter): Source[(String, InputTask), NotUsed]
-  def listCaseTasks(caseId: String): Source[(String, InputTask), NotUsed]
-  def listCaseTaskLogs(filter: Filter): Source[(String, InputLog), NotUsed]
-  def listCaseTaskLogs(caseId: String): Source[(String, InputLog), NotUsed]
-  def listAlerts(filter: Filter): Source[InputAlert, NotUsed]
-  def listAlertObservables(filter: Filter): Source[(String, InputObservable), NotUsed]
-  def listAlertObservables(alertId: String): Source[(String, InputObservable), NotUsed]
-  def listUsers(filter: Filter): Source[InputUser, NotUsed]
-  def listCustomFields(filter: Filter): Source[InputCustomField, NotUsed]
-  def listObservableTypes(filter: Filter): Source[InputObservableType, NotUsed]
-  def listProfiles(filter: Filter): Source[InputProfile, NotUsed]
-  def listImpactStatus(filter: Filter): Source[InputImpactStatus, NotUsed]
-  def listResolutionStatus(filter: Filter): Source[InputResolutionStatus, NotUsed]
-  def listCaseTemplate(filter: Filter): Source[InputCaseTemplate, NotUsed]
-  def listCaseTemplateTask(caseTemplateId: String): Source[(String, InputTask), NotUsed]
-  def listCaseTemplateTask(filter: Filter): Source[(String, InputTask), NotUsed]
-  def listJobs(caseId: String): Source[(String, InputJob), NotUsed]
-  def listJobs(filter: Filter): Source[(String, InputJob), NotUsed]
-  def listJobObservables(filter: Filter): Source[(String, InputObservable), NotUsed]
-  def listJobObservables(caseId: String): Source[(String, InputObservable), NotUsed]
-  def listAction(filter: Filter): Source[(String, InputAction), NotUsed]
-  def listAction(entityId: String): Source[(String, InputAction), NotUsed]
-  def listAudit(filter: Filter): Source[(String, InputAudit), NotUsed]
-  def listAudit(entityId: String, filter: Filter): Source[(String, InputAudit), NotUsed]
+  def listOrganisations(filter: Filter): Source[Try[InputOrganisation], NotUsed]
+  def countOrganisations(filter: Filter): Future[Long]
+  def listCases(filter: Filter): Source[Try[InputCase], NotUsed]
+  def countCases(filter: Filter): Future[Long]
+  def listCaseObservables(filter: Filter): Source[Try[(String, InputObservable)], NotUsed]
+  def countCaseObservables(filter: Filter): Future[Long]
+  def listCaseObservables(caseId: String): Source[Try[(String, InputObservable)], NotUsed]
+  def countCaseObservables(caseId: String): Future[Long]
+  def listCaseTasks(filter: Filter): Source[Try[(String, InputTask)], NotUsed]
+  def countCaseTasks(filter: Filter): Future[Long]
+  def listCaseTasks(caseId: String): Source[Try[(String, InputTask)], NotUsed]
+  def countCaseTasks(caseId: String): Future[Long]
+  def listCaseTaskLogs(filter: Filter): Source[Try[(String, InputLog)], NotUsed]
+  def countCaseTaskLogs(filter: Filter): Future[Long]
+  def listCaseTaskLogs(caseId: String): Source[Try[(String, InputLog)], NotUsed]
+  def countCaseTaskLogs(caseId: String): Future[Long]
+  def listAlerts(filter: Filter): Source[Try[InputAlert], NotUsed]
+  def countAlerts(filter: Filter): Future[Long]
+  def listAlertObservables(filter: Filter): Source[Try[(String, InputObservable)], NotUsed]
+  def countAlertObservables(filter: Filter): Future[Long]
+  def listAlertObservables(alertId: String): Source[Try[(String, InputObservable)], NotUsed]
+  def countAlertObservables(alertId: String): Future[Long]
+  def listUsers(filter: Filter): Source[Try[InputUser], NotUsed]
+  def countUsers(filter: Filter): Future[Long]
+  def listCustomFields(filter: Filter): Source[Try[InputCustomField], NotUsed]
+  def countCustomFields(filter: Filter): Future[Long]
+  def listObservableTypes(filter: Filter): Source[Try[InputObservableType], NotUsed]
+  def countObservableTypes(filter: Filter): Future[Long]
+  def listProfiles(filter: Filter): Source[Try[InputProfile], NotUsed]
+  def countProfiles(filter: Filter): Future[Long]
+  def listImpactStatus(filter: Filter): Source[Try[InputImpactStatus], NotUsed]
+  def countImpactStatus(filter: Filter): Future[Long]
+  def listResolutionStatus(filter: Filter): Source[Try[InputResolutionStatus], NotUsed]
+  def countResolutionStatus(filter: Filter): Future[Long]
+  def listCaseTemplate(filter: Filter): Source[Try[InputCaseTemplate], NotUsed]
+  def countCaseTemplate(filter: Filter): Future[Long]
+  def listCaseTemplateTask(caseTemplateId: String): Source[Try[(String, InputTask)], NotUsed]
+  def countCaseTemplateTask(caseTemplateId: String): Future[Long]
+  def listCaseTemplateTask(filter: Filter): Source[Try[(String, InputTask)], NotUsed]
+  def countCaseTemplateTask(filter: Filter): Future[Long]
+  def listJobs(caseId: String): Source[Try[(String, InputJob)], NotUsed]
+  def countJobs(caseId: String): Future[Long]
+  def listJobs(filter: Filter): Source[Try[(String, InputJob)], NotUsed]
+  def countJobs(filter: Filter): Future[Long]
+  def listJobObservables(filter: Filter): Source[Try[(String, InputObservable)], NotUsed]
+  def countJobObservables(filter: Filter): Future[Long]
+  def listJobObservables(caseId: String): Source[Try[(String, InputObservable)], NotUsed]
+  def countJobObservables(caseId: String): Future[Long]
+  def listAction(filter: Filter): Source[Try[(String, InputAction)], NotUsed]
+  def countAction(filter: Filter): Future[Long]
+  def listAction(entityId: String): Source[Try[(String, InputAction)], NotUsed]
+  def countAction(entityId: String): Future[Long]
+  def listAudit(filter: Filter): Source[Try[(String, InputAudit)], NotUsed]
+  def countAudit(filter: Filter): Future[Long]
+  def listAudit(entityId: String, filter: Filter): Source[Try[(String, InputAudit)], NotUsed]
+  def countAudit(entityId: String, filter: Filter): Future[Long]
 }
diff --git a/migration/src/main/scala/org/thp/thehive/migration/Migrate.scala b/migration/src/main/scala/org/thp/thehive/migration/Migrate.scala
@@ -13,6 +13,7 @@ import akka.actor.ActorSystem
 import akka.stream.Materializer
 import com.typesafe.config.{Config, ConfigFactory}
 import scopt.OParser
+import scala.concurrent.duration.DurationInt
 
 object Migrate extends App with MigrationOps {
   def getVersion: String = Option(getClass.getPackage.getImplementationVersion).getOrElse("SNAPSHOT")
@@ -90,9 +91,14 @@ object Migrate extends App with MigrationOps {
     val output = th4.Output(Configuration(config.getConfig("output").withFallback(config)))
     val filter = Filter.fromConfig(config.getConfig("input.filter"))
 
-    val process        = migrate(input, output, filter)
-    val migrationStats = Await.result(process, Duration.Inf)
+    val process = migrate(input, output, filter)
+    actorSystem.scheduler.scheduleAtFixedRate(1.seconds, 1.seconds) { () =>
+      logger.info(migrationStats.showStats())
+      migrationStats.flush()
+    }
+    Await.result(process, Duration.Inf)
     println("Migration finished")
+    migrationStats.flush()
     println(migrationStats)
     System.exit(0)
   }