#1340 Remove locks on unique constraint

TheHive-Project · May 22, 2020 · b63a15e · b63a15e
1 parent 4ebfaab
commit b63a15e
Show file tree

Hide file tree

Showing 6 changed files with 142 additions and 14 deletions.
diff --git a/thehive/app/org/thp/thehive/TheHiveModule.scala b/thehive/app/org/thp/thehive/TheHiveModule.scala
@@ -1,15 +1,17 @@
 package org.thp.thehive
 
-import play.api.libs.concurrent.AkkaGuiceSupport
+import akka.actor.{ActorRef, PoisonPill}
+import akka.cluster.singleton.{ClusterSingletonManager, ClusterSingletonManagerSettings}
 import com.google.inject.AbstractModule
 import net.codingwell.scalaguice.{ScalaModule, ScalaMultibinder}
 import org.thp.scalligraph.auth._
 import org.thp.scalligraph.janus.JanusDatabase
 import org.thp.scalligraph.models.{Database, Schema}
 import org.thp.scalligraph.services.{HadoopStorageSrv, S3StorageSrv}
-import org.thp.thehive.services.TOTPAuthSrvProvider
-import org.thp.thehive.services.notification.notifiers.{AppendToFileProvider, EmailerProvider, MattermostProvider, NotifierProvider, WebhookProvider}
+import org.thp.thehive.services.notification.notifiers._
 import org.thp.thehive.services.notification.triggers._
+import org.thp.thehive.services.{CaseDedupActor, CaseDedupActorProvider, DataDedupActor, DataDedupActorProvider, TOTPAuthSrvProvider}
+import play.api.libs.concurrent.AkkaGuiceSupport
 //import org.thp.scalligraph.orientdb.{OrientDatabase, OrientDatabaseStorageSrv}
 import org.thp.scalligraph.services.config.ConfigActor
 import org.thp.scalligraph.services.{DatabaseStorageSrv, LocalFileSystemStorageSrv, StorageSrv}
@@ -18,12 +20,11 @@ import org.thp.thehive.services.notification.NotificationActor
 import org.thp.thehive.services.{Connector, LocalKeyAuthProvider, LocalPasswordAuthProvider, LocalUserSrv}
 //import org.thp.scalligraph.neo4j.Neo4jDatabase
 //import org.thp.scalligraph.orientdb.OrientDatabase
-import play.api.routing.{Router => PlayRouter}
-import play.api.{Configuration, Environment, Logger}
-
 import org.thp.scalligraph.query.QueryExecutor
 import org.thp.thehive.controllers.v0.{TheHiveQueryExecutor => TheHiveQueryExecutorV0}
 import org.thp.thehive.controllers.v1.{TheHiveQueryExecutor => TheHiveQueryExecutorV1}
+import play.api.routing.{Router => PlayRouter}
+import play.api.{Configuration, Environment, Logger}
 
 class TheHiveModule(environment: Environment, configuration: Configuration) extends AbstractModule with ScalaModule with AkkaGuiceSupport {
   lazy val logger: Logger = Logger(getClass)
@@ -86,6 +87,36 @@ class TheHiveModule(environment: Environment, configuration: Configuration) exte
     bindActor[ConfigActor]("config-actor")
     bindActor[NotificationActor]("notification-actor")
 
+    bindActor[DataDedupActor](
+      "data-dedup-actor-singleton",
+      props =>
+        ClusterSingletonManager
+          .props(
+            singletonProps = props,
+            terminationMessage = PoisonPill,
+            settings = ClusterSingletonManagerSettings(configuration.get[Configuration]("akka.cluster.singleton").underlying)
+          )
+    )
+    bind[ActorRef]
+      .annotatedWithName("data-dedup-actor")
+      .toProvider[DataDedupActorProvider]
+      .asEagerSingleton()
+
+    bindActor[CaseDedupActor](
+      "case-dedup-actor-singleton",
+      props =>
+        ClusterSingletonManager
+          .props(
+            singletonProps = props,
+            terminationMessage = PoisonPill,
+            settings = ClusterSingletonManagerSettings(configuration.get[Configuration]("akka.cluster.singleton").underlying)
+          )
+    )
+    bind[ActorRef]
+      .annotatedWithName("case-dedup-actor")
+      .toProvider[CaseDedupActorProvider]
+      .asEagerSingleton()
+
     bind[SchemaUpdater].asEagerSingleton()
     bind[ClusterSetup].asEagerSingleton()
     ()

diff --git a/thehive/app/org/thp/thehive/models/Case.scala b/thehive/app/org/thp/thehive/models/Case.scala
@@ -60,7 +60,7 @@ case class CaseUser()
 case class CaseCaseTemplate()
 
 @VertexEntity
-@DefineIndex(IndexType.unique, "number")
+@DefineIndex(IndexType.tryUnique, "number")
 //@DefineIndex(IndexType.fulltext, "title")
 //@DefineIndex(IndexType.fulltext, "description")
 //@DefineIndex(IndexType.standard, "startDate")

diff --git a/thehive/app/org/thp/thehive/models/Observable.scala b/thehive/app/org/thp/thehive/models/Observable.scala
@@ -41,6 +41,6 @@ case class RichObservable(
   def sighted: Boolean           = observable.sighted
 }
 
-@DefineIndex(IndexType.unique, "data")
+@DefineIndex(IndexType.tryUnique, "data")
 @VertexEntity
 case class Data(data: String)
diff --git a/thehive/app/org/thp/thehive/models/SchemaUpdater.scala b/thehive/app/org/thp/thehive/models/SchemaUpdater.scala
@@ -2,7 +2,9 @@ package org.thp.thehive.models
 
 import gremlin.scala._
 import javax.inject.{Inject, Singleton}
+import org.janusgraph.core.schema.ConsistencyModifier
 import org.thp.scalligraph.auth.UserSrv
+import org.thp.scalligraph.janus.JanusDatabase
 import org.thp.scalligraph.models.{Database, IndexType, Operations}
 import org.thp.scalligraph.steps.StepsOps._
 import play.api.Logger
@@ -37,5 +39,27 @@ class SchemaUpdater @Inject() (theHiveSchema: TheHiveSchema, db: Database, userS
       Success(())
     }
     .addIndex("CustomField", IndexType.unique, "name")
+    .forVersion(4)
+    .dbOperation[JanusDatabase]("Remove locks") { db =>
+      def removePropertyLock(name: String) =
+        db.managementTransaction { mgmt =>
+          Try(mgmt.setConsistency(mgmt.getPropertyKey(name), ConsistencyModifier.DEFAULT))
+            .recover {
+              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")
+            }
+        }
+
+      removeIndexLock("CaseNumber")
+      removePropertyLock("number")
+      removeIndexLock("DataData")
+      removePropertyLock("data")
+    }
     .execute(db)(userSrv.getSystemAuthContext)
 }
diff --git a/thehive/app/org/thp/thehive/services/CaseSrv.scala b/thehive/app/org/thp/thehive/services/CaseSrv.scala
@@ -2,8 +2,10 @@ package org.thp.thehive.services
 
 import java.util.{List => JList, Set => JSet}
 
+import akka.actor.{ActorRef, ActorSystem}
+import akka.cluster.singleton.{ClusterSingletonProxy, ClusterSingletonProxySettings}
 import gremlin.scala._
-import javax.inject.{Inject, Singleton}
+import javax.inject.{Inject, Named, Provider, Singleton}
 import org.apache.tinkerpop.gremlin.process.traversal.{Order, Path, P => JP}
 import org.thp.scalligraph.auth.{AuthContext, Permission}
 import org.thp.scalligraph.controllers.FPathElem
@@ -18,6 +20,7 @@ import org.thp.thehive.models._
 import play.api.libs.json.{JsNull, JsObject, Json}
 
 import scala.collection.JavaConverters._
+import scala.concurrent.duration.{DurationInt, FiniteDuration}
 import scala.util.{Failure, Success, Try}
 
 @Singleton
@@ -32,7 +35,8 @@ class CaseSrv @Inject() (
     observableSrv: ObservableSrv,
     auditSrv: AuditSrv,
     resolutionStatusSrv: ResolutionStatusSrv,
-    impactStatusSrv: ImpactStatusSrv
+    impactStatusSrv: ImpactStatusSrv,
+    @Named("case-dedup-actor") caseDedupActor: ActorRef
 )(implicit db: Database)
     extends VertexSrv[Case, CaseSteps] {
 
@@ -44,6 +48,12 @@ class CaseSrv @Inject() (
   val caseCaseTemplateSrv     = new EdgeSrv[CaseCaseTemplate, Case, CaseTemplate]
   val mergedFromSrv           = new EdgeSrv[MergedFrom, Case, Case]
 
+  override def createEntity(e: Case)(implicit graph: Graph, authContext: AuthContext): Try[Case with Entity] =
+    super.createEntity(e).map { `case` =>
+      caseDedupActor ! DedupActor.EntityAdded
+      `case`
+    }
+
   def create(
       `case`: Case,
       user: Option[User with Entity],
@@ -595,3 +605,33 @@ class CaseSteps(raw: GremlinScala[Vertex])(implicit db: Database, graph: Graph)
 
   def alert: AlertSteps = new AlertSteps(raw.inTo[AlertCase])
 }
+
+class CaseDedupActor @Inject() (val db: Database, val service: CaseSrv) extends DedupActor[Case] {
+  override val min: FiniteDuration = 5.seconds
+  override val max: FiniteDuration = 10.seconds
+
+  override def resolve(entities: List[Case with Entity])(implicit graph: Graph): Try[Unit] = {
+    val nextNumber = service.nextCaseNumber
+    entities
+      .sorted(createdFirst)
+      .tail
+      .flatMap(service.get(_).raw.headOption())
+      .zipWithIndex
+      .foreach {
+        case (vertex, index) =>
+          db.setSingleProperty(vertex, "number", nextNumber + index, UniMapping.int)
+      }
+    Success(())
+  }
+}
+
+class CaseDedupActorProvider @Inject() (system: ActorSystem, @Named("case-dedup-actor-singleton") CaseDedupActorSingleton: ActorRef)
+    extends Provider[ActorRef] {
+  override def get(): ActorRef =
+    system.actorOf(
+      ClusterSingletonProxy.props(
+        singletonManagerPath = CaseDedupActorSingleton.path.toStringWithoutAddress,
+        settings = ClusterSingletonProxySettings(system)
+      )
+    )
+}
diff --git a/thehive/app/org/thp/thehive/services/DataSrv.scala b/thehive/app/org/thp/thehive/services/DataSrv.scala
@@ -2,10 +2,10 @@ package org.thp.thehive.services
 
 import java.lang.{Long => JLong}
 
-import scala.util.{Success, Try}
-
+import akka.actor.{ActorRef, ActorSystem}
+import akka.cluster.singleton.{ClusterSingletonProxy, ClusterSingletonProxySettings}
 import gremlin.scala.{Graph, GremlinScala, P, Vertex}
-import javax.inject.{Inject, Singleton}
+import javax.inject.{Inject, Named, Provider, Singleton}
 import org.apache.tinkerpop.gremlin.structure.T
 import org.thp.scalligraph.EntitySteps
 import org.thp.scalligraph.auth.AuthContext
@@ -15,10 +15,19 @@ import org.thp.scalligraph.steps.StepsOps._
 import org.thp.scalligraph.steps.{Traversal, VertexSteps}
 import org.thp.thehive.models._
 
+import scala.concurrent.duration.{DurationInt, FiniteDuration}
+import scala.util.{Success, Try}
+
 @Singleton
-class DataSrv @Inject() ()(implicit db: Database) extends VertexSrv[Data, DataSteps] {
+class DataSrv @Inject() (@Named("data-dedup-actor") dataDedupActor: ActorRef)(implicit db: Database) extends VertexSrv[Data, DataSteps] {
   override def steps(raw: GremlinScala[Vertex])(implicit graph: Graph): DataSteps = new DataSteps(raw)
 
+  override def createEntity(e: Data)(implicit graph: Graph, authContext: AuthContext): Try[Data with Entity] =
+    super.createEntity(e).map { data =>
+      dataDedupActor ! DedupActor.EntityAdded
+      data
+    }
+
   def create(e: Data)(implicit graph: Graph, authContext: AuthContext): Try[Data with Entity] =
     initSteps
       .getByData(e.data)
@@ -49,3 +58,27 @@ class DataSteps(raw: GremlinScala[Vertex])(implicit db: Database, graph: Graph)
 
   def useCount: Traversal[JLong, JLong] = Traversal(raw.inTo[ObservableData].count())
 }
+
+class DataDedupActor @Inject() (val db: Database, val service: DataSrv) extends DedupActor[Data] {
+  override val min: FiniteDuration = 10.seconds
+  override val max: FiniteDuration = 1.minute
+
+  override def resolve(entities: List[Data with Entity])(implicit graph: Graph): Try[Unit] = entities match {
+    case head :: tail =>
+      tail.foreach(copyEdge(_, head))
+      tail.foreach(service.get(_).remove())
+      Success(())
+    case _ => Success(())
+  }
+}
+
+class DataDedupActorProvider @Inject() (system: ActorSystem, @Named("data-dedup-actor-singleton") DataDedupActorSingleton: ActorRef)
+    extends Provider[ActorRef] {
+  override def get(): ActorRef =
+    system.actorOf(
+      ClusterSingletonProxy.props(
+        singletonManagerPath = DataDedupActorSingleton.path.toStringWithoutAddress,
+        settings = ClusterSingletonProxySettings(system)
+      )
+    )
+}