From 21bba420224e14d3d55d7017ba8b423f70cf2a75 Mon Sep 17 00:00:00 2001
From: To-om <toom@thehive-project.org>
Date: Tue, 5 Jan 2021 14:22:48 +0100
Subject: [PATCH] #1731 Update database schema

---
 .../models/TheHiveSchemaDefinition.scala      | 49 +++++++++++++++++++
 1 file changed, 49 insertions(+)

diff --git a/thehive/app/org/thp/thehive/models/TheHiveSchemaDefinition.scala b/thehive/app/org/thp/thehive/models/TheHiveSchemaDefinition.scala
index 3e16940379..c9fe95e6f4 100644
--- a/thehive/app/org/thp/thehive/models/TheHiveSchemaDefinition.scala
+++ b/thehive/app/org/thp/thehive/models/TheHiveSchemaDefinition.scala
@@ -14,6 +14,7 @@ import org.thp.scalligraph.janus.JanusDatabase
 import org.thp.scalligraph.models._
 import org.thp.scalligraph.traversal.Graph
 import org.thp.scalligraph.traversal.TraversalOps._
+import org.thp.thehive.services.LocalUserSrv
 import play.api.Logger
 
 import scala.collection.JavaConverters._
@@ -83,11 +84,57 @@ class TheHiveSchemaDefinition @Inject() extends Schema with UpdatableSchema {
         .iterate()
       Success(())
     }
+    //=====[release 4.0.2]=====
     .addProperty[Boolean]("ShareTask", "actionRequired")
     .updateGraph("Add actionRequire property", "Share") { traversal =>
       traversal.outE[ShareTask].raw.property("actionRequired", false).iterate()
       Success(())
     }
+    //=====[release 4.0.3]=====
+    .addProperty[String]("Alert", "organisationId")
+    .updateGraph("Add organisation data in alerts", "Alert") { traversal =>
+      traversal
+        .project(_.by.by(_.out("AlertOrganisation")._id))
+        .toIterator
+        .foreach {
+          case (vertex, organisationId) =>
+            vertex.property("organisationId", organisationId.value)
+        }
+      Success(())
+    }
+    .addProperty[Seq[String]]("Case", "organisationIds")
+    .updateGraph("Add organisation data in cases", "Case") { traversal =>
+      traversal
+        .project(_.by.by(_.in("ShareCase").in("OrganisationShare")._id.fold))
+        .toIterator
+        .foreach {
+          case (vertex, organisationIds) =>
+            organisationIds.foreach(id => vertex.property(Cardinality.list, "organisationIds", id.value))
+        }
+      Success(())
+    }
+    .addProperty[Seq[String]]("Observable", "organisationIds")
+    .updateGraph("Add organisation data in observables", "Observable") { traversal =>
+      traversal
+        .project(
+          _.by
+            .by(
+              _.coalesceIdent(
+                _.optional(_.in("ReportObservable").in("ObservableJob")).in("ShareObservable").in("OrganisationShare"),
+                _.in("AlertObservable").out("AlertOrganisation")
+              )
+                ._id
+                .fold
+            )
+        )
+        .toIterator
+        .foreach {
+          case (vertex, organisationIds) =>
+            organisationIds.foreach(id => vertex.property(Cardinality.list, "organisationIds", id.value))
+        }
+      Success(())
+    }
+    .rebuildIndexes
 
   val reflectionClasses = new Reflections(
     new ConfigurationBuilder()
@@ -116,4 +163,6 @@ class TheHiveSchemaDefinition @Inject() extends Schema with UpdatableSchema {
   }.flatten
 
   override def init(db: Database)(implicit graph: Graph, authContext: AuthContext): Try[Unit] = Success(())
+
+  override val authContext: AuthContext = LocalUserSrv.getSystemAuthContext
 }