From c5b134347ef1cb800fbca7f3118df980ea53a54f Mon Sep 17 00:00:00 2001
From: To-om <toom@thehive-project.org>
Date: Tue, 9 Mar 2021 10:04:27 +0100
Subject: [PATCH] #1785 Enforce database initialisation

---
 .../migration/th4/JanusDatabaseProvider.scala | 19 ++++++++++++++++--
 .../thp/thehive/migration/th4/Output.scala    | 20 +------------------
 2 files changed, 18 insertions(+), 21 deletions(-)

diff --git a/migration/src/main/scala/org/thp/thehive/migration/th4/JanusDatabaseProvider.scala b/migration/src/main/scala/org/thp/thehive/migration/th4/JanusDatabaseProvider.scala
index 1927e0cc1c..ef5e9b898b 100644
--- a/migration/src/main/scala/org/thp/thehive/migration/th4/JanusDatabaseProvider.scala
+++ b/migration/src/main/scala/org/thp/thehive/migration/th4/JanusDatabaseProvider.scala
@@ -1,25 +1,40 @@
 package org.thp.thehive.migration.th4
 
 import akka.actor.ActorSystem
+import org.janusgraph.core.JanusGraph
 import org.thp.scalligraph.SingleInstance
 import org.thp.scalligraph.janus.JanusDatabase
 import org.thp.scalligraph.models.{Database, UpdatableSchema}
 import play.api.Configuration
 
 import javax.inject.{Inject, Provider, Singleton}
+import scala.collection.JavaConverters._
 import scala.collection.immutable
 
 @Singleton
 class JanusDatabaseProvider @Inject() (configuration: Configuration, system: ActorSystem, schemas: immutable.Set[UpdatableSchema])
     extends Provider[Database] {
+
+  def dropOtherConnections(db: JanusGraph): Unit = {
+    val mgmt = db.openManagement()
+    mgmt
+      .getOpenInstances
+      .asScala
+      .filterNot(_.endsWith("(current)"))
+      .foreach(mgmt.forceCloseInstance)
+    mgmt.commit()
+  }
+
   override lazy val get: Database = {
+    val janusDatabase = JanusDatabase.openDatabase(configuration, system)
+    dropOtherConnections(janusDatabase)
     val db = new JanusDatabase(
-      JanusDatabase.openDatabase(configuration, system),
+      janusDatabase,
       configuration,
       system,
       new SingleInstance(true)
     )
-    schemas.foreach(schema => db.createSchemaFrom(schema)(schema.authContext))
+    schemas.toTry(schema => schema.update(db)).get
     db
   }
 }
diff --git a/migration/src/main/scala/org/thp/thehive/migration/th4/Output.scala b/migration/src/main/scala/org/thp/thehive/migration/th4/Output.scala
index 46fd86af85..12ce91c110 100644
--- a/migration/src/main/scala/org/thp/thehive/migration/th4/Output.scala
+++ b/migration/src/main/scala/org/thp/thehive/migration/th4/Output.scala
@@ -242,25 +242,7 @@ class Output @Inject() (
                      | ${tags.size} tags""".stripMargin)
   }
 
-  def startMigration(): Try[Unit] =
-//    db match {
-//      case jdb: JanusDatabase => jdb.dropOtherConnections.recover { case error => logger.error(s"Fail to remove other connection", error) }
-//      case _                  =>
-//    }
-    if (db.version("thehive") == 0)
-      db.createSchemaFrom(theHiveSchema)(LocalUserSrv.getSystemAuthContext)
-        .flatMap(_ => db.setVersion(theHiveSchema.name, theHiveSchema.operations.lastVersion))
-        .flatMap(_ => db.createSchemaFrom(cortexSchema)(LocalUserSrv.getSystemAuthContext))
-        .flatMap(_ => db.setVersion(cortexSchema.name, cortexSchema.operations.lastVersion))
-        .map(_ => retrieveExistingData())
-    else
-      theHiveSchema
-        .update(db)
-        .flatMap(_ => cortexSchema.update(db))
-        .map { _ =>
-          retrieveExistingData()
-          db.rebuildIndexes()
-        }
+  def startMigration(): Try[Unit] = Success(retrieveExistingData())
 
   def endMigration(): Try[Unit] = {
     db.addSchemaIndexes(theHiveSchema)