#1404 Improve integrity checks

TheHive-Project · Jul 5, 2020 · 5855f88 · 5855f88
1 parent 0c4bf82
commit 5855f88
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 22 deletions.
diff --git a/thehive/app/org/thp/thehive/services/CaseTemplateSrv.scala b/thehive/app/org/thp/thehive/services/CaseTemplateSrv.scala
@@ -3,9 +3,10 @@ package org.thp.thehive.services
 import java.util.{Collection => JCollection, List => JList, Map => JMap}
 
 import akka.actor.ActorRef
-import gremlin.scala.{__, By, Element, Graph, GremlinScala, Key, P, StepLabel, Vertex}
+import gremlin.scala.{__, By, Element, Graph, GremlinScala, Key, P, Vertex}
 import javax.inject.{Inject, Named}
 import org.apache.tinkerpop.gremlin.process.traversal.{Path, Scope}
+import org.apache.tinkerpop.gremlin.structure.T
 import org.thp.scalligraph.auth.{AuthContext, Permission}
 import org.thp.scalligraph.models.{Database, Entity}
 import org.thp.scalligraph.query.PropertyUpdater
@@ -251,24 +252,28 @@ class CaseTemplateSteps(raw: GremlinScala[Vertex])(implicit @Named("with-thehive
     new CustomFieldValueSteps(raw.outToE[CaseTemplateCustomField])
 }
 
-class CaseTemplateIntegrityCheckOps @Inject() (@Named("with-thehive-schema") val db: Database, val service: CaseTemplateSrv)
-    extends IntegrityCheckOps[CaseTemplate] {
-  override def getDuplicates[A](property: String): List[List[CaseTemplate with Entity]] =
+class CaseTemplateIntegrityCheckOps @Inject() (
+    @Named("with-thehive-schema") val db: Database,
+    val service: CaseTemplateSrv,
+    organisationSrv: OrganisationSrv
+) extends IntegrityCheckOps[CaseTemplate] {
+  override def duplicateEntities: List[List[CaseTemplate with Entity]] =
     db.roTransaction { implicit graph =>
-      val ctLabel  = StepLabel()
-      val orgLabel = StepLabel()
-      service
+      organisationSrv
         .initSteps
-        .as(ctLabel)
-        .organisation
-        .as(orgLabel)
         .raw
-        .group(By(__.select(ctLabel.name, orgLabel.name).by("name").by()))
-        .unfold[JMap.Entry[Any, JCollection[Vertex]]]()
-        .selectValues
-        .where(_.count(Scope.local).is(P.gt(1)))
+        .traversal
+        .flatMap(
+          __.in("CaseTemplateOrganisation")
+            .group(By(Key[String]("name")), By(T.id))
+            .unfold[JMap.Entry[String, JCollection[Any]]]()
+            .selectValues
+            .where(_.count(Scope.local).is(P.gt(1)))
+            .traversal
+        )
+        .asScala
+        .map(ids => service.getByIds(ids.asScala.map(_.toString).toSeq: _*).toList)
         .toList
-        .map(_.asScala.toList.map(service.model.toDomain(_)(db)))
     }
 
   override def resolve(entities: List[CaseTemplate with Entity])(implicit graph: Graph): Try[Unit] = entities match {

diff --git a/thehive/app/org/thp/thehive/services/TagSrv.scala b/thehive/app/org/thp/thehive/services/TagSrv.scala
@@ -65,7 +65,7 @@ class TagSrv @Inject() (appConfig: ApplicationConfig, @Named("integrity-check-ac
   override def exists(e: Tag)(implicit graph: Graph): Boolean = initSteps.getByName(e.namespace, e.predicate, e.value).exists()
 }
 
-class TagSteps(raw: GremlinScala[Vertex])(implicit @Named("with-thehive-schema") db: Database, graph: Graph) extends VertexSteps[Tag](raw) {
+class TagSteps(raw: GremlinScala[Vertex])(implicit db: Database, graph: Graph) extends VertexSteps[Tag](raw) {
   override def newInstance(newRaw: GremlinScala[Vertex]): TagSteps = new TagSteps(newRaw)
   override def newInstance(): TagSteps                             = new TagSteps(raw.clone())
 
@@ -84,6 +84,7 @@ class TagSteps(raw: GremlinScala[Vertex])(implicit @Named("with-thehive-schema")
 }
 
 class TagIntegrityCheckOps @Inject() (@Named("with-thehive-schema") val db: Database, val service: TagSrv) extends IntegrityCheckOps[Tag] {
+
   override def resolve(entities: List[Tag with Entity])(implicit graph: Graph): Try[Unit] = entities match {
     case head :: tail =>
       tail.foreach(copyEdge(_, head))

diff --git a/thehive/app/org/thp/thehive/services/UserSrv.scala b/thehive/app/org/thp/thehive/services/UserSrv.scala
@@ -338,12 +338,7 @@ class UserIntegrityCheckOps @Inject() (
   }
 
   override def check(): Unit = {
-    duplicateEntities
-      .foreach { entities =>
-        db.tryTransaction { implicit graph =>
-          resolve(entities)
-        }
-      }
+    super.check()
     db.tryTransaction { implicit graph =>
       duplicateInEdges[TaskUser](service.initSteps.raw).flatMap(firstCreatedElement(_)).foreach(e => removeEdges(e._2))
       duplicateInEdges[CaseUser](service.initSteps.raw).flatMap(firstCreatedElement(_)).foreach(e => removeEdges(e._2))