#1404 Check integrity at startup

TheHive-Project · Jun 28, 2020 · a0cfe0e · a0cfe0e
1 parent 5d335ae
commit a0cfe0e
Show file tree

Hide file tree

Showing 4 changed files with 55 additions and 19 deletions.
diff --git a/ScalliGraph b/ScalliGraph
diff --git a/...rg/thp/thehive/models/TheHiveSchema.scala → ...hive/models/TheHiveSchemaDefinition.scala b/...rg/thp/thehive/models/TheHiveSchema.scala → ...hive/models/TheHiveSchemaDefinition.scala
@@ -41,7 +41,7 @@ class TheHiveSchemaDefinition @Inject() (injector: Injector) extends Schema with
       }
       Success(())
     }
-    .addIndex("CustomField", IndexType.unique, "name")
+    .noop // .addIndex("CustomField", IndexType.unique, "name")
     .dbOperation[JanusDatabase]("Remove locks") { db =>
       def removePropertyLock(name: String) =
         db.managementTransaction { mgmt =>
@@ -50,21 +50,22 @@ class TheHiveSchemaDefinition @Inject() (injector: Injector) extends Schema with
               case error => logger.warn(s"Unable to remove lock on property $name: $error")
             }
         }
-      def removeIndexLock(name: String) =
-        db.managementTransaction { mgmt =>
-          Try(mgmt.setConsistency(mgmt.getGraphIndex(name), ConsistencyModifier.DEFAULT))
-            .recover {
-              case error => logger.warn(s"Unable to remove lock on index $name: $error")
-            }
-        }
+      // def removeIndexLock(name: String): Try[Unit] =
+      //   db.managementTransaction { mgmt =>
+      //     Try(mgmt.setConsistency(mgmt.getGraphIndex(name), ConsistencyModifier.DEFAULT))
+      //       .recover {
+      //         case error => logger.warn(s"Unable to remove lock on index $name: $error")
+      //       }
+      //   }
 
-      removeIndexLock("CaseNumber")
+      // removeIndexLock("CaseNumber")
       removePropertyLock("number")
-      removeIndexLock("DataData")
+      // removeIndexLock("DataData")
       removePropertyLock("data")
     }
-    .addIndex("Tag", IndexType.tryUnique, "namespace", "predicate", "value")
-    .addIndex("Audit", IndexType.basic, "requestId", "mainAction")
+    .noop // .addIndex("Tag", IndexType.unique, "namespace", "predicate", "value")
+    .noop // .addIndex("Audit", IndexType.basic, "requestId", "mainAction")
+    .rebuildIndexes
 
   val reflectionClasses = new Reflections(
     new ConfigurationBuilder()

diff --git a/thehive/test/org/thp/thehive/DatabaseBuilder.scala b/thehive/test/org/thp/thehive/DatabaseBuilder.scala
@@ -3,13 +3,13 @@ package org.thp.thehive
 import java.io.File
 
 import gremlin.scala.{KeyValue => _, _}
-import javax.inject.{Inject, Named, Singleton}
+import javax.inject.{Inject, Singleton}
 import org.scalactic.Or
 import org.thp.scalligraph.RichOption
 import org.thp.scalligraph.auth.AuthContext
 import org.thp.scalligraph.controllers._
 import org.thp.scalligraph.models.{Database, Entity, Schema}
-import org.thp.scalligraph.services.{EdgeSrv, VertexSrv}
+import org.thp.scalligraph.services.{EdgeSrv, GenIntegrityCheckOps, IntegrityCheckOps, VertexSrv}
 import org.thp.thehive.models._
 import org.thp.thehive.services._
 import play.api.Logger
@@ -42,18 +42,25 @@ class DatabaseBuilder @Inject() (
     alertSrv: AlertSrv,
     attachmentSrv: AttachmentSrv,
     dashboardSrv: DashboardSrv,
-    pageSrv: PageSrv
+    pageSrv: PageSrv,
+    integrityChecks: Set[GenIntegrityCheckOps]
 ) {
 
   lazy val logger: Logger = Logger(getClass)
 
-  def build()(implicit @Named("with-thehive-schema") db: Database, authContext: AuthContext): Try[Unit] = {
+  def build()(implicit db: Database, authContext: AuthContext): Try[Unit] = {
 
     lazy val logger: Logger = Logger(getClass)
     logger.info("Initialize database schema")
     db.createSchemaFrom(schema)
       .flatMap(_ => db.addSchemaIndexes(schema))
       .flatMap { _ =>
+        integrityChecks.foreach { check =>
+          db.tryTransaction { implicit graph =>
+            Success(check.initialCheck())
+          }
+          ()
+        }
         db.tryTransaction { implicit graph =>
           val idMap =
             createVertex(caseSrv, FieldsParser[Case]) ++

diff --git a/thehive/test/org/thp/thehive/TestAppBuilder.scala b/thehive/test/org/thp/thehive/TestAppBuilder.scala
@@ -7,12 +7,27 @@ import javax.inject.{Inject, Provider, Singleton}
 import org.apache.commons.io.FileUtils
 import org.thp.scalligraph.auth._
 import org.thp.scalligraph.models.{Database, Schema}
-import org.thp.scalligraph.services.{LocalFileSystemStorageSrv, StorageSrv}
+import org.thp.scalligraph.services.{GenIntegrityCheckOps, LocalFileSystemStorageSrv, StorageSrv}
 import org.thp.scalligraph.{janus, AppBuilder}
 import org.thp.thehive.models.TheHiveSchemaDefinition
 import org.thp.thehive.services.notification.notifiers.{AppendToFileProvider, EmailerProvider, NotifierProvider}
 import org.thp.thehive.services.notification.triggers._
-import org.thp.thehive.services.{LocalKeyAuthProvider, LocalPasswordAuthProvider, LocalUserSrv}
+import org.thp.thehive.services.{
+  CaseIntegrityCheckOps,
+  CaseTemplateIntegrityCheckOps,
+  CustomFieldIntegrityCheckOps,
+  DataIntegrityCheckOps,
+  ImpactStatusIntegrityCheckOps,
+  LocalKeyAuthProvider,
+  LocalPasswordAuthProvider,
+  LocalUserSrv,
+  ObservableTypeIntegrityCheckOps,
+  OrganisationIntegrityCheckOps,
+  ProfileIntegrityCheckOps,
+  ResolutionStatusIntegrityCheckOps,
+  TagIntegrityCheckOps,
+  UserIntegrityCheckOps
+}
 
 object TestAppBuilderLock
 
@@ -33,6 +48,19 @@ trait TestAppBuilder {
       .multiBind[TriggerProvider](classOf[TaskAssignedProvider])
       .multiBind[TriggerProvider](classOf[AlertCreatedProvider])
       .bindToProvider[AuthSrv, MultiAuthSrvProvider]
+      .multiBind[GenIntegrityCheckOps](
+        classOf[ProfileIntegrityCheckOps],
+        classOf[OrganisationIntegrityCheckOps],
+        classOf[TagIntegrityCheckOps],
+        classOf[UserIntegrityCheckOps],
+        classOf[ImpactStatusIntegrityCheckOps],
+        classOf[ResolutionStatusIntegrityCheckOps],
+        classOf[ObservableTypeIntegrityCheckOps],
+        classOf[CustomFieldIntegrityCheckOps],
+        classOf[CaseTemplateIntegrityCheckOps],
+        classOf[DataIntegrityCheckOps],
+        classOf[CaseIntegrityCheckOps]
+      )
       .bindActor[DummyActor]("config-actor")
       .bindActor[DummyActor]("notification-actor")
       .bindActor[DummyActor]("integrity-check-actor")
+1 −0		core/src/main/scala/org/thp/scalligraph/models/Database.scala
+1 −1		core/src/main/scala/org/thp/scalligraph/models/Model.scala
+10 −0		core/src/main/scala/org/thp/scalligraph/models/Operation.scala
+5 −1		core/src/main/scala/org/thp/scalligraph/services/IntegrityCheckOps.scala
+12 −12		database/janusgraph/src/main/scala/org/thp/scalligraph/janus/JanusDatabase.scala