#1387 Prevent duplicated observable type, impact status and resolutio…

…n status
TheHive-Project · Jun 29, 2020 · 3a56149 · 3a56149
1 parent f9a6d69
commit 3a56149
Show file tree

Hide file tree

Showing 5 changed files with 22 additions and 8 deletions.
diff --git a/thehive/app/org/thp/thehive/models/Case.scala b/thehive/app/org/thp/thehive/models/Case.scala
@@ -14,6 +14,7 @@ object CaseStatus extends Enumeration {
 }
 
 @VertexEntity
+@DefineIndex(IndexType.unique, "value")
 case class ResolutionStatus(value: String) {
   require(!value.isEmpty, "ResolutionStatus can't be empty")
 }
@@ -32,6 +33,7 @@ object ResolutionStatus {
 case class CaseResolutionStatus()
 
 @VertexEntity
+@DefineIndex(IndexType.unique, "value")
 case class ImpactStatus(value: String) {
   require(!value.isEmpty, "ImpactStatus can't be empty")
 }

diff --git a/thehive/app/org/thp/thehive/models/ObservableType.scala b/thehive/app/org/thp/thehive/models/ObservableType.scala
@@ -1,11 +1,13 @@
 package org.thp.thehive.models
 
+import org.thp.scalligraph.models.{DefineIndex, IndexType}
 import org.thp.scalligraph.{EdgeEntity, VertexEntity}
 
 @EdgeEntity[Observable, ObservableType]
 case class ObservableObservableType()
 
 @VertexEntity
+@DefineIndex(IndexType.unique, "name")
 case class ObservableType(name: String, isAttachment: Boolean)
 
 object ObservableType {

diff --git a/thehive/app/org/thp/thehive/services/ImpactStatusSrv.scala b/thehive/app/org/thp/thehive/services/ImpactStatusSrv.scala
@@ -3,15 +3,15 @@ package org.thp.thehive.services
 import akka.actor.ActorRef
 import gremlin.scala._
 import javax.inject.{Inject, Named, Singleton}
-import org.thp.scalligraph.EntitySteps
+import org.thp.scalligraph.{CreateError, EntitySteps}
 import org.thp.scalligraph.auth.AuthContext
 import org.thp.scalligraph.models.{Database, Entity}
 import org.thp.scalligraph.services.{IntegrityCheckOps, VertexSrv}
 import org.thp.scalligraph.steps.StepsOps._
 import org.thp.scalligraph.steps.VertexSteps
 import org.thp.thehive.models.ImpactStatus
 
-import scala.util.{Success, Try}
+import scala.util.{Failure, Success, Try}
 
 @Singleton
 class ImpactStatusSrv @Inject() (@Named("integrity-check-actor") integrityCheckActor: ActorRef)(
@@ -29,7 +29,11 @@ class ImpactStatusSrv @Inject() (@Named("integrity-check-actor") integrityCheckA
     super.createEntity(e)
   }
 
-  def create(impactStatus: ImpactStatus)(implicit graph: Graph, authContext: AuthContext): Try[ImpactStatus with Entity] = createEntity(impactStatus)
+  def create(impactStatus: ImpactStatus)(implicit graph: Graph, authContext: AuthContext): Try[ImpactStatus with Entity] =
+    if (exists(impactStatus))
+      Failure(CreateError(s"Impact status ${impactStatus.value} already exists"))
+    else
+      createEntity(impactStatus)
 
   override def exists(e: ImpactStatus)(implicit graph: Graph): Boolean = initSteps.getByName(e.value).exists()
 }

diff --git a/thehive/app/org/thp/thehive/services/ObservableTypeSrv.scala b/thehive/app/org/thp/thehive/services/ObservableTypeSrv.scala
@@ -8,7 +8,7 @@ import org.thp.scalligraph.models.{Database, Entity}
 import org.thp.scalligraph.services._
 import org.thp.scalligraph.steps.StepsOps._
 import org.thp.scalligraph.steps.VertexSteps
-import org.thp.scalligraph.{BadRequestError, EntitySteps}
+import org.thp.scalligraph.{BadRequestError, CreateError, EntitySteps}
 import org.thp.thehive.models._
 
 import scala.util.{Failure, Success, Try}
@@ -33,7 +33,10 @@ class ObservableTypeSrv @Inject() (@Named("integrity-check-actor") integrityChec
   }
 
   def create(observableType: ObservableType)(implicit graph: Graph, authContext: AuthContext): Try[ObservableType with Entity] =
-    createEntity(observableType)
+    if (exists(observableType))
+      Failure(CreateError(s"Observable type ${observableType.name} already exists"))
+    else
+      createEntity(observableType)
 
   def remove(idOrName: String)(implicit graph: Graph): Try[Unit] =
     if (useCount(idOrName) == 0) Success(get(idOrName).remove())

diff --git a/thehive/app/org/thp/thehive/services/ResolutionStatusSrv.scala b/thehive/app/org/thp/thehive/services/ResolutionStatusSrv.scala
@@ -3,15 +3,15 @@ package org.thp.thehive.services
 import akka.actor.ActorRef
 import gremlin.scala._
 import javax.inject.{Inject, Named, Singleton}
-import org.thp.scalligraph.EntitySteps
+import org.thp.scalligraph.{CreateError, EntitySteps}
 import org.thp.scalligraph.auth.AuthContext
 import org.thp.scalligraph.models.{Database, Entity}
 import org.thp.scalligraph.services.{IntegrityCheckOps, VertexSrv}
 import org.thp.scalligraph.steps.StepsOps._
 import org.thp.scalligraph.steps.VertexSteps
 import org.thp.thehive.models.ResolutionStatus
 
-import scala.util.{Success, Try}
+import scala.util.{Failure, Success, Try}
 
 @Singleton
 class ResolutionStatusSrv @Inject() (@Named("integrity-check-actor") integrityCheckActor: ActorRef)(
@@ -30,7 +30,10 @@ class ResolutionStatusSrv @Inject() (@Named("integrity-check-actor") integrityCh
   }
 
   def create(resolutionStatus: ResolutionStatus)(implicit graph: Graph, authContext: AuthContext): Try[ResolutionStatus with Entity] =
-    createEntity(resolutionStatus)
+    if (exists(resolutionStatus))
+      Failure(CreateError(s"Resolution status ${resolutionStatus.value} already exists"))
+    else
+      createEntity(resolutionStatus)
 
   override def exists(e: ResolutionStatus)(implicit graph: Graph): Boolean = initSteps.getByName(e.value).exists()
 }