#1731 Fix duplication checks

TheHive-Project · Jan 22, 2021 · 03632c3 · 03632c3
1 parent 01524dc
commit 03632c3
Show file tree

Hide file tree

Showing 5 changed files with 45 additions and 84 deletions.
diff --git a/thehive/app/org/thp/thehive/models/Case.scala b/thehive/app/org/thp/thehive/models/Case.scala
@@ -107,12 +107,12 @@ case class Case(
     status: CaseStatus.Value,
     summary: Option[String],
     tags: Seq[String],
+    impactStatus: Option[String] = None,
+    resolutionStatus: Option[String] = None,
     /* filled by the service */
     assignee: Option[String] = None,
     number: Int = 0,
     organisationIds: Seq[EntityId] = Nil,
-    impactStatus: Option[String] = None,
-    resolutionStatus: Option[String] = None,
     caseTemplate: Option[String] = None
 )
 

diff --git a/thehive/app/org/thp/thehive/services/AlertSrv.scala b/thehive/app/org/thp/thehive/services/AlertSrv.scala
@@ -129,41 +129,22 @@ class AlertSrv @Inject() (
   def createObservable(alert: Alert with Entity, observable: Observable, data: String)(implicit
       graph: Graph,
       authContext: AuthContext
-  ): Try[RichObservable] = {
-    val alreadyExists = observableSrv
-      .startTraversal
-      .has(_.relatedId, alert._id)
-      .has(_.data, data)
-      .exists
-    if (alreadyExists)
-      Failure(CreateError("Observable already exists"))
-    else
-      for {
-        createdObservable <- observableSrv.create(observable.copy(organisationIds = Seq(organisationSrv.currentId), relatedId = alert._id), data)
-        _                 <- alertObservableSrv.create(AlertObservable(), alert, createdObservable.observable)
-        _                 <- auditSrv.observableInAlert.create(createdObservable.observable, alert, createdObservable.toJson)
-      } yield createdObservable
-  }
+  ): Try[RichObservable] =
+    for {
+      createdObservable <- observableSrv.create(observable.copy(organisationIds = Seq(organisationSrv.currentId), relatedId = alert._id), data)
+      _                 <- alertObservableSrv.create(AlertObservable(), alert, createdObservable.observable)
+      _                 <- auditSrv.observableInAlert.create(createdObservable.observable, alert, createdObservable.toJson)
+    } yield createdObservable
 
   def createObservable(alert: Alert with Entity, observable: Observable, attachment: Attachment with Entity)(implicit
       graph: Graph,
       authContext: AuthContext
-  ): Try[RichObservable] = {
-    val alreadyExists = observableSrv
-      .startTraversal
-      .has(_.relatedId, alert._id)
-      .has(_.attachmentId, attachment.attachmentId)
-      .exists
-    if (alreadyExists)
-      Failure(CreateError("Observable already exists"))
-    else
-      for {
-        createdObservable <-
-          observableSrv.create(observable.copy(organisationIds = Seq(organisationSrv.currentId), relatedId = alert._id), attachment)
-        _ <- alertObservableSrv.create(AlertObservable(), alert, createdObservable.observable)
-        _ <- auditSrv.observableInAlert.create(createdObservable.observable, alert, createdObservable.toJson)
-      } yield createdObservable
-  }
+  ): Try[RichObservable] =
+    for {
+      createdObservable <- observableSrv.create(observable.copy(organisationIds = Seq(organisationSrv.currentId), relatedId = alert._id), attachment)
+      _                 <- alertObservableSrv.create(AlertObservable(), alert, createdObservable.observable)
+      _                 <- auditSrv.observableInAlert.create(createdObservable.observable, alert, createdObservable.toJson)
+    } yield createdObservable
 
   def createObservable(alert: Alert with Entity, observable: Observable, file: FFile)(implicit
       graph: Graph,

diff --git a/thehive/app/org/thp/thehive/services/CaseSrv.scala b/thehive/app/org/thp/thehive/services/CaseSrv.scala
@@ -16,7 +16,6 @@ import org.thp.thehive.controllers.v1.Conversion._
 import org.thp.thehive.dto.v1.InputCustomFieldValue
 import org.thp.thehive.models._
 import org.thp.thehive.services.CaseOps._
-import org.thp.thehive.services.CaseTemplateOps._
 import org.thp.thehive.services.CustomFieldOps._
 import org.thp.thehive.services.DataOps._
 import org.thp.thehive.services.ObservableOps._
@@ -35,7 +34,6 @@ class CaseSrv @Inject() (
     tagSrv: TagSrv,
     customFieldSrv: CustomFieldSrv,
     organisationSrv: OrganisationSrv,
-    caseTemplateSrv: CaseTemplateSrv,
     profileSrv: ProfileSrv,
     shareSrv: ShareSrv,
     taskSrv: TaskSrv,
@@ -70,9 +68,10 @@ class CaseSrv @Inject() (
       caseTemplate: Option[RichCaseTemplate],
       additionalTasks: Seq[Task]
   )(implicit graph: Graph, authContext: AuthContext): Try[RichCase] = {
-    val tags       = (`case`.tags ++ caseTemplate.fold[Seq[String]](Nil)(_.tags)).distinct
     val caseNumber = if (`case`.number == 0) nextCaseNumber else `case`.number
+    val tagNames   = (`case`.tags ++ caseTemplate.fold[Seq[String]](Nil)(_.tags)).distinct
     for {
+      tags <- tagNames.toTry(tagSrv.getOrCreate)
       createdCase <- createEntity(
         `case`.copy(
           number = caseNumber,
@@ -81,14 +80,14 @@ class CaseSrv @Inject() (
           caseTemplate = caseTemplate.map(_.name),
           impactStatus = None,
           resolutionStatus = None,
-          tags = tags
+          tags = tagNames
         )
       )
-      _            <- assignee.map(u => caseUserSrv.create(CaseUser(), createdCase, u)).flip
-      _            <- shareSrv.shareCase(owner = true, createdCase, organisation, profileSrv.orgAdmin)
-      _            <- caseTemplate.map(ct => caseCaseTemplateSrv.create(CaseCaseTemplate(), createdCase, ct.caseTemplate)).flip
-      _            <- caseTemplate.fold(additionalTasks)(_.tasks.map(_.task) ++ additionalTasks).toTry(task => createTask(createdCase, task))
-      caseTemplate <- `case`.caseTemplate.map(caseTemplateSrv.getByName(_).richCaseTemplate.getOrFail("CaseTemplate")).flip
+      _ <- assignee.map(u => caseUserSrv.create(CaseUser(), createdCase, u)).flip
+      _ <- shareSrv.shareCase(owner = true, createdCase, organisation, profileSrv.orgAdmin)
+      _ <- caseTemplate.map(ct => caseCaseTemplateSrv.create(CaseCaseTemplate(), createdCase, ct.caseTemplate)).flip
+      _ <- caseTemplate.fold(additionalTasks)(_.tasks.map(_.task) ++ additionalTasks).toTry(task => createTask(createdCase, task))
+      _ <- tags.toTry(caseTagSrv.create(CaseTag(), createdCase, _))
 
       caseTemplateCf =
         caseTemplate
@@ -177,41 +176,20 @@ class CaseSrv @Inject() (
   def createObservable(`case`: Case with Entity, observable: Observable, data: String)(implicit
       graph: Graph,
       authContext: AuthContext
-  ): Try[RichObservable] = {
-    val alreadyExists = observableSrv
-      .startTraversal
-      .has(_.organisationIds, organisationSrv.currentId)
-      .has(_.relatedId, `case`._id)
-      .has(_.data, data)
-      .exists
-    if (alreadyExists)
-      Failure(CreateError("Observable already exists"))
-    else
-      for {
-        createdObservable <- observableSrv.create(observable.copy(organisationIds = Seq(organisationSrv.currentId), relatedId = `case`._id), data)
-        _                 <- shareSrv.shareObservable(createdObservable, `case`, organisationSrv.currentId)
-      } yield createdObservable
-  }
+  ): Try[RichObservable] =
+    for {
+      createdObservable <- observableSrv.create(observable.copy(organisationIds = Seq(organisationSrv.currentId), relatedId = `case`._id), data)
+      _                 <- shareSrv.shareObservable(createdObservable, `case`, organisationSrv.currentId)
+    } yield createdObservable
 
   def createObservable(`case`: Case with Entity, observable: Observable, attachment: Attachment with Entity)(implicit
       graph: Graph,
       authContext: AuthContext
-  ): Try[RichObservable] = {
-    val alreadyExists = observableSrv
-      .startTraversal
-      .has(_.organisationIds, organisationSrv.currentId)
-      .has(_.relatedId, `case`._id)
-      .has(_.attachmentId, attachment.attachmentId)
-      .exists
-    if (alreadyExists)
-      Failure(CreateError("Observable already exists"))
-    else
-      for {
-        createdObservable <-
-          observableSrv.create(observable.copy(organisationIds = Seq(organisationSrv.currentId), relatedId = `case`._id), attachment)
-        _ <- shareSrv.shareObservable(createdObservable, `case`, organisationSrv.currentId)
-      } yield createdObservable
-  }
+  ): Try[RichObservable] =
+    for {
+      createdObservable <- observableSrv.create(observable.copy(organisationIds = Seq(organisationSrv.currentId), relatedId = `case`._id), attachment)
+      _                 <- shareSrv.shareObservable(createdObservable, `case`, organisationSrv.currentId)
+    } yield createdObservable
 
   def createObservable(`case`: Case with Entity, observable: Observable, file: FFile)(implicit
       graph: Graph,

diff --git a/thehive/app/org/thp/thehive/services/ObservableSrv.scala b/thehive/app/org/thp/thehive/services/ObservableSrv.scala
@@ -10,7 +10,7 @@ import org.thp.scalligraph.services._
 import org.thp.scalligraph.traversal.TraversalOps._
 import org.thp.scalligraph.traversal.{Converter, Graph, StepLabel, Traversal}
 import org.thp.scalligraph.utils.Hash
-import org.thp.scalligraph.{BadRequestError, CreateError, EntityId, EntityIdOrName, RichSeq}
+import org.thp.scalligraph.{BadRequestError, CreateError, EntityId, EntityIdOrName, EntityName, RichSeq}
 import org.thp.thehive.models._
 import org.thp.thehive.services.ObservableOps._
 import org.thp.thehive.services.OrganisationOps._
@@ -60,17 +60,19 @@ class ObservableSrv @Inject() (
       .has(_.dataType, observable.dataType)
       .filterOnAttachmentId(attachment.attachmentId)
       .exists
-    if (alreadyExists)
+    if (alreadyExists) Failure(CreateError("Observable already exists"))
+    else
       for {
-        observableType <- observableTypeSrv.getOrFail(EntityIdOrName(observable.dataType))
+        observableType <- observableTypeSrv.getOrFail(EntityName(observable.dataType))
         _ <-
           if (!observableType.isAttachment) Failure(BadRequestError("A text observable doesn't accept attachment"))
           else Success(())
+        tags              <- observable.tags.toTry(tagSrv.getOrCreate)
         createdObservable <- createEntity(observable.copy(data = None))
         _                 <- observableObservableType.create(ObservableObservableType(), createdObservable, observableType)
         _                 <- observableAttachmentSrv.create(ObservableAttachment(), createdObservable, attachment)
+        _                 <- tags.toTry(observableTagSrv.create(ObservableTag(), createdObservable, _))
       } yield RichObservable(createdObservable, Some(attachment), None, Nil)
-    else Failure(CreateError("Observable already exists"))
   }
 
   def create(
@@ -86,18 +88,20 @@ class ObservableSrv @Inject() (
       .has(_.data, observable.data.get)
       .has(_.dataType, observable.dataType)
       .exists
-    if (alreadyExists)
+    if (alreadyExists) Failure(CreateError("Observable already exists"))
+    else
       for {
-        observableType <- observableTypeSrv.getOrFail(EntityIdOrName(observable.dataType))
+        observableType <- observableTypeSrv.getOrFail(EntityName(observable.dataType))
         _ <-
           if (observableType.isAttachment) Failure(BadRequestError("A attachment observable doesn't accept string value"))
           else Success(())
+        tags              <- observable.tags.toTry(tagSrv.getOrCreate)
         data              <- dataSrv.create(Data(dataValue))
         createdObservable <- createEntity(observable.copy(data = Some(dataValue)))
         _                 <- observableObservableType.create(ObservableObservableType(), createdObservable, observableType)
         _                 <- observableDataSrv.create(ObservableData(), createdObservable, data)
+        _                 <- tags.toTry(observableTagSrv.create(ObservableTag(), createdObservable, _))
       } yield RichObservable(createdObservable, None, None, Nil)
-    else Failure(CreateError("Observable already exists"))
   }
 
   def addTags(observable: Observable with Entity, tags: Set[String])(implicit graph: Graph, authContext: AuthContext): Try[Seq[Tag with Entity]] = {

diff --git a/thehive/test/org/thp/thehive/services/CaseSrvTest.scala b/thehive/test/org/thp/thehive/services/CaseSrvTest.scala
@@ -1,21 +1,19 @@
 package org.thp.thehive.services
 
-import java.util.Date
-
 import org.specs2.matcher.Matcher
+import org.thp.scalligraph.EntityName
 import org.thp.scalligraph.auth.AuthContext
 import org.thp.scalligraph.controllers.FPathElem
 import org.thp.scalligraph.models._
 import org.thp.scalligraph.query.PropertyUpdater
 import org.thp.scalligraph.traversal.TraversalOps._
-import org.thp.scalligraph.{CreateError, EntityName}
 import org.thp.thehive.TestAppBuilder
 import org.thp.thehive.models._
 import org.thp.thehive.services.CaseOps._
-import org.thp.thehive.services.ObservableOps._
 import play.api.libs.json.Json
 import play.api.test.PlaySpecification
 
+import java.util.Date
 import scala.util.Success
 
 class CaseSrvTest extends PlaySpecification with TestAppBuilder {