#2033 Add integrity check for tasks, logs and observables

thehive/app/org/thp/thehive/TheHiveModule.scala

@@ -102,6 +102,9 @@ class TheHiveModule(environment: Environment, configuration: Configuration) exte
     integrityCheckOpsBindings.addBinding.to[DataIntegrityCheckOps]
     integrityCheckOpsBindings.addBinding.to[CaseIntegrityCheckOps]
     integrityCheckOpsBindings.addBinding.to[AlertIntegrityCheckOps]
+    integrityCheckOpsBindings.addBinding.to[TaskIntegrityCheckOps]
+    integrityCheckOpsBindings.addBinding.to[ObservableIntegrityCheckOps]
+    integrityCheckOpsBindings.addBinding.to[LogIntegrityCheckOps]
     bind[ActorRef].annotatedWithName("integrity-check-actor").toProvider[IntegrityCheckActorProvider]
 
     bind[ActorRef].annotatedWithName("flow-actor").toProvider[FlowActorProvider]

thehive/app/org/thp/thehive/models/Case.scala

@@ -173,7 +173,9 @@ object RichCase {
       assignee: Option[String],
       customFields: Seq[RichCustomField],
       userPermissions: Set[Permission],
-      organisationIds: Set[EntityId]
+      organisationIds: Set[EntityId],
+      caseTemplate: Option[String] = None,
+      owningOrganisation: EntityId = EntityId.empty
   ): RichCase = {
     val `case`: Case with Entity =
       new Case(
@@ -192,7 +194,9 @@ object RichCase {
         tags = tags,
         assignee = assignee,
         impactStatus = impactStatus,
-        resolutionStatus = resolutionStatus
+        resolutionStatus = resolutionStatus,
+        caseTemplate = caseTemplate,
+        owningOrganisation = owningOrganisation
       ) with Entity {
         override val _id: EntityId              = __id
         override val _label: String             = "Case"

thehive/app/org/thp/thehive/services/CaseSrv.scala

@@ -712,9 +712,7 @@ class CaseIntegrityCheckOps @Inject() (
     Success(())
   }
 
-  override def globalCheck(): Map[String, Int] = {
-    implicit val authContext: AuthContext = LocalUserSrv.getSystemAuthContext
-
+  override def globalCheck(): Map[String, Int] =
     db.tryTransaction { implicit graph =>
       Try {
         service
@@ -724,23 +722,28 @@ class CaseIntegrityCheckOps @Inject() (
               .by(_.organisations._id.fold)
               .by(_.assignee.value(_.login).fold)
               .by(_.caseTemplate.value(_.name).fold)
+              .by(_.origin._id.fold)
           )
           .toIterator
           .map {
-            case (case0, organisationIds, assigneeIds, caseTemplateNames) =>
+            case (case0, organisationIds, assigneeIds, caseTemplateNames, owningOrganisationIds) =>
+              val fixOwningOrg: LinkRemover =
+                (caseId, orgId) => service.get(caseId).shares.filter(_.organisation.get(orgId._id)).update(_.owner, false).iterate()
+
               val assigneeStats = singleOptionLink[User, String]("assignee", userSrv.getByName(_).head, _.login)(_.outEdge[CaseUser])
                 .check(case0, case0.assignee, assigneeIds)
-              val orgStats = multiIdLink[Organisation]("organisationIds", organisationSrv)(_.remove)
-                .check(case0, case0.organisationIds.toSeq, organisationIds)
+              val orgStats = multiIdLink[Organisation]("organisationIds", organisationSrv)(_.remove) // FIXME => Seq => Set
+                .check(case0, case0.organisationIds, organisationIds)
               val templateStats =
                 singleOptionLink[CaseTemplate, String]("caseTemplate", caseTemplateSrv.getByName(_).head, _.name)(_.outEdge[CaseCaseTemplate])
                   .check(case0, case0.caseTemplate, caseTemplateNames)
+              val owningOrgStats = singleIdLink[Organisation]("owningOrganisation", organisationSrv)(_ => fixOwningOrg, _.remove)
+                .check(case0, case0.owningOrganisation, owningOrganisationIds)
 
-              assigneeStats <+> orgStats <+> templateStats
+              assigneeStats <+> orgStats <+> templateStats <+> owningOrgStats
           }
           .reduceOption(_ <+> _)
           .getOrElse(Map.empty)
       }
     }.getOrElse(Map("globalFailure" -> 1))
-  }
 }

thehive/app/org/thp/thehive/services/CaseTemplateSrv.scala

@@ -237,9 +237,7 @@ object CaseTemplateOps {
         }
         .getOrElse(traversal.empty)
 
-    def hasCustomField(customFieldSrv: CustomFieldSrv, customField: EntityIdOrName): Traversal.V[CaseTemplate] = {
-      val cfFilter = (t: Traversal.V[CustomField]) => customField.fold(id => t.hasId(id), name => t.has(_.name, name))
-
+    def hasCustomField(customFieldSrv: CustomFieldSrv, customField: EntityIdOrName): Traversal.V[CaseTemplate] =
       customFieldSrv
         .get(customField)(traversal.graph)
         .value(_.`type`)
@@ -252,11 +250,8 @@ object CaseTemplateOps {
           case CustomFieldType.string  => traversal.filter(_.customFields.has(_.stringValue).inV.v[CustomField].get(customField))
         }
         .getOrElse(traversal.empty)
-    }
-
-    def hasNotCustomField(customFieldSrv: CustomFieldSrv, customField: EntityIdOrName): Traversal.V[CaseTemplate] = {
-      val cfFilter = (t: Traversal.V[CustomField]) => customField.fold(id => t.hasId(id), name => t.has(_.name, name))
 
+    def hasNotCustomField(customFieldSrv: CustomFieldSrv, customField: EntityIdOrName): Traversal.V[CaseTemplate] =
       customFieldSrv
         .get(customField)(traversal.graph)
         .value(_.`type`)
@@ -269,7 +264,6 @@ object CaseTemplateOps {
           case CustomFieldType.string  => traversal.filterNot(_.customFields.has(_.stringValue).inV.v[CustomField].get(customField))
         }
         .getOrElse(traversal.empty)
-    }
   }
 
   implicit class CaseTemplateCustomFieldsOpsDefs(traversal: Traversal.E[CaseTemplateCustomField]) extends CustomFieldValueOpsDefs(traversal)

thehive/app/org/thp/thehive/services/IntegrityCheckActor.scala

@@ -219,6 +219,7 @@ class IntegrityCheckActor() extends Actor {
         }
       }
     case GlobalCheckResult(name, stats) =>
+      logger.info(s"End of $name global check: $stats")
       states.get(name).foreach { state =>
         context.become(receive(states + (name -> state.copy(globalStats = state.globalStats + stats))))
       }

thehive/app/org/thp/thehive/services/LogSrv.scala

@@ -112,9 +112,7 @@ object LogOps {
 class LogIntegrityCheckOps @Inject() (val db: Database, val service: LogSrv, taskSrv: TaskSrv) extends IntegrityCheckOps[Log] {
   override def resolve(entities: Seq[Log with Entity])(implicit graph: Graph): Try[Unit] = Success(())
 
-  override def globalCheck(): Map[String, Int] = {
-    implicit val authContext: AuthContext = LocalUserSrv.getSystemAuthContext
-
+  override def globalCheck(): Map[String, Int] =
     db.tryTransaction { implicit graph =>
       Try {
         service
@@ -133,5 +131,4 @@ class LogIntegrityCheckOps @Inject() (val db: Database, val service: LogSrv, tas
           .getOrElse(Map.empty)
       }
     }.getOrElse(Map("globalFailure" -> 1))
-  }
 }

thehive/app/org/thp/thehive/services/ObservableSrv.scala

@@ -4,7 +4,7 @@ import org.apache.tinkerpop.gremlin.process.traversal.P
 import org.apache.tinkerpop.gremlin.structure.Vertex
 import org.thp.scalligraph.auth.{AuthContext, Permission}
 import org.thp.scalligraph.controllers.FFile
-import org.thp.scalligraph.models.{Database, Entity}
+import org.thp.scalligraph.models.{Database, Entity, UMapping}
 import org.thp.scalligraph.query.PropertyUpdater
 import org.thp.scalligraph.services._
 import org.thp.scalligraph.traversal.Converter.Identity
@@ -237,7 +237,10 @@ object ObservableOps {
         .domainMap(profile => profile.permissions & authContext.permissions)
 
     def organisations: Traversal.V[Organisation] =
-      traversal.coalesceIdent(_.in[ShareObservable].in[OrganisationShare], _.in[AlertObservable].out[AlertOrganisation]).v[Organisation]
+      traversal
+        .unionFlat(identity, _.in("ReportObservable").in("ObservableJob").v[Observable])
+        .unionFlat(_.shares.organisation, _.alert.organisation)
+//      traversal.coalesceIdent(_.in[ShareObservable].in[OrganisationShare], _.in[AlertObservable].out[AlertOrganisation]).v[Organisation]
 
     def origin: Traversal.V[Organisation] = shares.has(_.owner, true).organisation
 
@@ -355,18 +358,43 @@ object ObservableOps {
   }
 }
 
-class ObservableIntegrityCheckOps @Inject() (val db: Database, val service: ObservableSrv) extends IntegrityCheckOps[Observable] {
+class ObservableIntegrityCheckOps @Inject() (val db: Database, val service: ObservableSrv, organisationSrv: OrganisationSrv)
+    extends IntegrityCheckOps[Observable] {
   override def resolve(entities: Seq[Observable with Entity])(implicit graph: Graph): Try[Unit] = Success(())
 
   override def globalCheck(): Map[String, Int] =
     db.tryTransaction { implicit graph =>
       Try {
-        val orphanIds = service.startTraversal.filterNot(_.or(_.shares, _.alert, _.in("ReportObservable")))._id.toSeq
-        if (orphanIds.nonEmpty) {
-          logger.warn(s"Found ${orphanIds.length} observables orphan(s) (${orphanIds.mkString(",")})")
-          service.getByIds(orphanIds: _*).remove()
-        }
-        Map("orphans" -> orphanIds.size)
+        service
+          .startTraversal
+          .project(
+            _.by
+              .by(_.organisations._id.fold)
+              .by(_.unionFlat(_.`case`._id, _.alert._id, _.in("ReportObservable")._id).fold)
+          )
+          .toIterator
+          .map {
+            case (observable, organisationIds, relatedIds) =>
+              val orgStats = multiIdLink[Organisation]("organisationIds", organisationSrv)(_.remove)
+                .check(observable, observable.organisationIds, organisationIds)
+
+              val removeOrphan: OrphanStrategy[Observable, EntityId] = { (a, entity) =>
+                service.get(entity).remove()
+                Map("Observable-relatedId-removeOrphan" -> 1)
+              }
+              val relatedStats = new SingleLinkChecker[Product, EntityId, EntityId](
+                orphanStrategy = removeOrphan,
+                setField = (entity, link) => UMapping.entityId.setProperty(service.get(entity), "relatedId", link._id).iterate(),
+                entitySelector = _ => EntitySelector.firstCreatedEntity,
+                removeLink = (_, _) => (),
+                getLink = id => graph.VV(id).entity.head,
+                Some(_)
+              ).check(observable, observable.relatedId, relatedIds)
+
+              orgStats <+> relatedStats
+          }
+          .reduceOption(_ <+> _)
+          .getOrElse(Map.empty)
       }
     }.getOrElse(Map("globalFailure" -> 1))
 }

thehive/app/org/thp/thehive/services/TaskSrv.scala

@@ -3,7 +3,7 @@ package org.thp.thehive.services
 import org.apache.tinkerpop.gremlin.process.traversal.P
 import org.apache.tinkerpop.gremlin.structure.Vertex
 import org.thp.scalligraph.auth.{AuthContext, Permission}
-import org.thp.scalligraph.models.{Entity, Model}
+import org.thp.scalligraph.models.{Database, Entity, Model, UMapping}
 import org.thp.scalligraph.query.PropertyUpdater
 import org.thp.scalligraph.services._
 import org.thp.scalligraph.traversal.Converter.Identity
@@ -12,6 +12,7 @@ import org.thp.scalligraph.traversal.{Converter, Graph, Traversal}
 import org.thp.scalligraph.utils.FunctionalCondition._
 import org.thp.scalligraph.{EntityId, EntityIdOrName}
 import org.thp.thehive.models.{TaskStatus, _}
+import org.thp.thehive.services.CaseTemplateOps._
 import org.thp.thehive.services.OrganisationOps._
 import org.thp.thehive.services.ShareOps._
 import org.thp.thehive.services.TaskOps._
@@ -20,7 +21,7 @@ import play.api.libs.json.{JsNull, JsObject, Json}
 import java.lang.{Boolean => JBoolean}
 import java.util.{Date, Map => JMap}
 import javax.inject.{Inject, Provider, Singleton}
-import scala.util.{Failure, Try}
+import scala.util.{Failure, Success, Try}
 
 @Singleton
 class TaskSrv @Inject() (
@@ -236,3 +237,44 @@ object TaskOps {
       traversal.in[ShareTask].filter(_.in[OrganisationShare].v[Organisation].get(organisation)).v[Share]
   }
 }
+
+class TaskIntegrityCheckOps @Inject() (val db: Database, val service: TaskSrv, organisationSrv: OrganisationSrv) extends IntegrityCheckOps[Task] {
+  override def resolve(entities: Seq[Task with Entity])(implicit graph: Graph): Try[Unit] = Success(())
+
+  override def globalCheck(): Map[String, Int] =
+    db.tryTransaction { implicit graph =>
+      Try {
+        service
+          .startTraversal
+          .project(
+            _.by
+              .by(_.unionFlat(_.`case`._id, _.caseTemplate._id).fold)
+              .by(_.unionFlat(_.organisations._id, _.caseTemplate.organisation._id).fold)
+          )
+          .toIterator
+          .map {
+            case (task, relatedIds, organisationIds) =>
+              val orgStats = multiIdLink[Organisation]("organisationIds", organisationSrv)(_.remove)
+                .check(task, task.organisationIds, organisationIds)
+
+              val removeOrphan: OrphanStrategy[Task, EntityId] = { (a, entity) =>
+                service.get(entity).remove()
+                Map("Task-relatedId-removeOrphan" -> 1)
+              }
+
+              val relatedStats = new SingleLinkChecker[Product, EntityId, EntityId](
+                orphanStrategy = removeOrphan,
+                setField = (entity, link) => UMapping.entityId.setProperty(service.get(entity), "relatedId", link._id).iterate(),
+                entitySelector = _ => EntitySelector.firstCreatedEntity,
+                removeLink = (_, _) => (),
+                getLink = id => graph.VV(id).entity.head,
+                Some(_)
+              ).check(task, task.relatedId, relatedIds)
+
+              orgStats <+> relatedStats
+          }
+          .reduceOption(_ <+> _)
+          .getOrElse(Map.empty)
+      }
+    }.getOrElse(Map("globalFailure" -> 1))
+}

thehive/conf/reference.conf

@@ -197,6 +197,21 @@ integrityCheck {
     interval: 30 minutes
     globalInterval: 6 hours
   }
+  Task {
+    initialDelay: 5 minute
+    interval: 30 minutes
+    globalInterval: 6 hours
+  }
+  Log {
+    initialDelay: 5 minute
+    interval: 30 minutes
+    globalInterval: 6 hours
+  }
+  Observable {
+    initialDelay: 5 minute
+    interval: 30 minutes
+    globalInterval: 6 hours
+  }
 }