#1731 Use index in queries

cortex/connector/src/main/scala/org/thp/thehive/connector/cortex/services/ActionOperationSrv.scala

@@ -1,7 +1,5 @@
 package org.thp.thehive.connector.cortex.services
 
-import java.util.Date
-import javax.inject.Inject
 import org.thp.scalligraph.auth.AuthContext
 import org.thp.scalligraph.models.Entity
 import org.thp.scalligraph.traversal.Graph
@@ -11,10 +9,11 @@ import org.thp.thehive.connector.cortex.models._
 import org.thp.thehive.controllers.v0.Conversion._
 import org.thp.thehive.dto.v0.InputTask
 import org.thp.thehive.models._
-import org.thp.thehive.services.CaseOps._
 import org.thp.thehive.services._
 import play.api.Logger
 
+import java.util.Date
+import javax.inject.Inject
 import scala.util.{Failure, Success, Try}
 
 class ActionOperationSrv @Inject() (
@@ -24,9 +23,7 @@ class ActionOperationSrv @Inject() (
     alertSrv: AlertSrv,
     logSrv: LogSrv,
     organisationSrv: OrganisationSrv,
-    observableTypeSrv: ObservableTypeSrv,
-    userSrv: UserSrv,
-    shareSrv: ShareSrv
+    userSrv: UserSrv
 ) {
   private[ActionOperationSrv] lazy val logger: Logger = Logger(getClass)
 
@@ -62,10 +59,8 @@ class ActionOperationSrv @Inject() (
 
       case CreateTask(title, description) =>
         for {
-          case0        <- relatedCase.fold[Try[Case with Entity]](Failure(InternalError("Unable to apply action CreateTask without case")))(Success(_))
-          createdTask  <- taskSrv.create(InputTask(title = title, description = Some(description)).toTask, None)
-          organisation <- organisationSrv.getOrFail(authContext.organisation)
-          _            <- shareSrv.shareTask(createdTask, case0, organisation)
+          case0 <- relatedCase.fold[Try[Case with Entity]](Failure(InternalError("Unable to apply action CreateTask without case")))(Success(_))
+          _     <- caseSrv.createTask(case0, InputTask(title = title, description = Some(description)).toTask)
         } yield updateOperation(operation)
 
       case AddCustomFields(name, _, value) =>
@@ -92,34 +87,32 @@ class ActionOperationSrv @Inject() (
           _ <- logSrv.create(Log(content, new Date(), deleted = false), t, None)
         } yield updateOperation(operation)
 
-      case AddArtifactToCase(_, dataType, dataMessage) =>
+      case AddArtifactToCase(data, dataType, message) =>
         for {
           c            <- relatedCase.fold[Try[Case with Entity]](Failure(InternalError("Unable to apply action AddArtifactToCase without case")))(Success(_))
-          obsType      <- observableTypeSrv.getOrFail(EntityIdOrName(dataType))
           organisation <- organisationSrv.getOrFail(authContext.organisation)
-          richObservable <- observableSrv.create(
+          _ <- caseSrv.createObservable(
+            c,
             Observable(
-              Some(dataMessage),
-              2,
+              message = Some(message),
+              tlp = 2,
               ioc = false,
               sighted = false,
               ignoreSimilarity = None,
-              organisationIds = Seq(organisation._id),
-              relatedId = c._id
+              dataType = dataType,
+              tags = Nil,
+              relatedId = c._id,
+              organisationIds = Seq(organisation._id)
             ),
-            obsType,
-            dataMessage,
-            Set.empty[String],
-            Nil
+            data
           )
-          _ <- caseSrv.addObservable(c, richObservable)
         } yield updateOperation(operation)
 
       case AssignCase(owner) =>
         for {
           c <- relatedCase.fold[Try[Case with Entity]](Failure(InternalError("Unable to apply action AssignCase without case")))(Success(_))
           u <- userSrv.get(EntityIdOrName(owner)).getOrFail("User")
-          _ <- Try(caseSrv.startTraversal.getEntity(c).unassign())
+          _ <- caseSrv.unassign(c)
           _ <- caseSrv.assign(c, u)
         } yield updateOperation(operation)
 

cortex/connector/src/main/scala/org/thp/thehive/connector/cortex/services/AnalyzerSrv.scala

@@ -68,10 +68,10 @@ class AnalyzerSrv @Inject() (connector: Connector, serviceHelper: ServiceHelper,
   def getAnalyzerByName(analyzerName: String, organisation: EntityIdOrName): Future[Map[CortexWorker, Seq[String]]] =
     searchAnalyzers(Json.obj("query" -> Json.obj("_field" -> "name", "_value" -> analyzerName)), organisation)
 
-  def searchAnalyzers(query: JsObject)(implicit authContext: AuthContext): Future[Map[OutputWorker, Seq[String]]] =
+  def searchAnalyzers(query: JsObject)(implicit authContext: AuthContext): Future[Map[CortexWorker, Seq[String]]] =
     searchAnalyzers(query, authContext.organisation)
 
-  def searchAnalyzers(query: JsObject, organisation: EntityIdOrName): Future[Map[OutputWorker, Seq[String]]] =
+  def searchAnalyzers(query: JsObject, organisation: EntityIdOrName): Future[Map[CortexWorker, Seq[String]]] =
     Future
       .traverse(serviceHelper.availableCortexClients(connector.clients, organisation)) { client =>
         client

cortex/connector/src/main/scala/org/thp/thehive/connector/cortex/services/Conversion.scala

@@ -22,16 +22,21 @@ object Conversion {
 
   implicit class CortexOutputArtifactOps(artifact: OutputArtifact) {
 
-    def toObservable(relatedId: EntityId, organisationIds: EntityId*): Observable =
+    def toObservable(
+        relatedId: EntityId,
+        organisations: Seq[EntityId]
+    ): Observable =
       artifact
         .into[Observable]
         .withFieldComputed(_.message, _.message)
         .withFieldComputed(_.tlp, _.tlp)
         .withFieldConst(_.ioc, false)
         .withFieldConst(_.sighted, false)
         .withFieldConst(_.ignoreSimilarity, None)
-        .withFieldConst(_.organisationIds, organisationIds)
+        .withFieldConst(_.data, None)
+        .withFieldComputed(_.tags, _.tags.toSeq)
         .withFieldConst(_.relatedId, relatedId)
+        .withFieldConst(_.organisationIds, organisations)
         .transform
   }
 

cortex/connector/src/main/scala/org/thp/thehive/connector/cortex/services/JobSrv.scala

@@ -1,15 +1,11 @@
 package org.thp.thehive.connector.cortex.services
 
-import java.nio.file.Files
-import java.util.{Date, Map => JMap}
 import akka.Done
 import akka.actor._
 import akka.stream.Materializer
 import akka.stream.scaladsl.FileIO
 import com.google.inject.name.Named
 import io.scalaland.chimney.dsl._
-
-import javax.inject.{Inject, Singleton}
 import org.apache.tinkerpop.gremlin.process.traversal.P
 import org.thp.cortex.client.CortexClient
 import org.thp.cortex.dto.v0.{InputArtifact, OutputArtifact, Attachment => CortexAttachment, JobStatus => CortexJobStatus, OutputJob => CortexJob}
@@ -32,6 +28,9 @@ import org.thp.thehive.services.OrganisationOps._
 import org.thp.thehive.services.{AttachmentSrv, ObservableSrv, ObservableTypeSrv, ReportTagSrv}
 import play.api.libs.json.Json
 
+import java.nio.file.Files
+import java.util.{Date, Map => JMap}
+import javax.inject.{Inject, Singleton}
 import scala.concurrent.{ExecutionContext, Future}
 import scala.util.{Failure, Success, Try}
 
@@ -76,15 +75,15 @@ class JobSrv @Inject() (
       analyzer <- cortexClient.getAnalyzer(workerId).recoverWith {
         case _ => cortexClient.getAnalyzerByName(workerId)
       } // if get analyzer using cortex2 API fails, try using legacy API
-      cortexArtifact <- (observable.attachment, observable.data) match {
-        case (None, Some(data)) =>
+      cortexArtifact <- observable.dataOrAttachment match {
+        case Left(data) =>
           Future.successful(
-            InputArtifact(observable.tlp, `case`.pap, observable.`type`.name, `case`.number.toString, Some(data.data), None)
+            InputArtifact(observable.tlp, `case`.pap, observable.dataType, `case`.number.toString, Some(data), None)
           )
-        case (Some(a), None) =>
+        case Right(a) =>
           val attachment = CortexAttachment(a.name, a.size, a.contentType, attachmentSrv.source(a))
           Future.successful(
-            InputArtifact(observable.tlp, `case`.pap, observable.`type`.name, `case`.number.toString, None, Some(attachment))
+            InputArtifact(observable.tlp, `case`.pap, observable.dataType, `case`.number.toString, None, Some(attachment))
           )
         case _ => Future.failed(new Exception(s"Invalid Observable data for ${observable.observable._id}"))
       }
@@ -207,21 +206,15 @@ class JobSrv @Inject() (
     Future
       .traverse(artifacts) { artifact =>
         db.tryTransaction(graph => observableTypeSrv.getOrFail(EntityIdOrName(artifact.dataType))(graph)) match {
-          case Success(attachmentType) if attachmentType.isAttachment => importCortexAttachment(job, artifact, attachmentType, cortexClient)
-          case Success(dataType) =>
+          case Success(attachmentType) if attachmentType.isAttachment => importCortexAttachment(job, artifact, cortexClient)
+          case _: Success[_] =>
             Future
               .fromTry {
                 db.tryTransaction { implicit graph =>
                   for {
                     origObs <- get(job).observable.getOrFail("Observable")
-                    obs <- observableSrv.create(
-                      artifact.toObservable(job._id, origObs.organisationIds: _*),
-                      dataType,
-                      artifact.data.get,
-                      artifact.tags,
-                      Nil
-                    )
-                    _ <- addObservable(job, obs.observable)
+                    obs     <- observableSrv.create(artifact.toObservable(job._id, origObs.organisationIds), artifact.data.get)
+                    _       <- addObservable(job, obs.observable)
                   } yield ()
                 }
               }
@@ -248,7 +241,6 @@ class JobSrv @Inject() (
   private def importCortexAttachment(
       job: Job with Entity,
       artifact: OutputArtifact,
-      attachmentType: ObservableType with Entity,
       cortexClient: CortexClient
   )(implicit
       authContext: AuthContext
@@ -266,10 +258,8 @@ class JobSrv @Inject() (
               for {
                 origObs           <- get(job).observable.getOrFail("Observable")
                 createdAttachment <- attachmentSrv.create(fFile)
-                richObservable <-
-                  observableSrv
-                    .create(artifact.toObservable(job._id, origObs.organisationIds: _*), attachmentType, createdAttachment, artifact.tags, Nil)
-                _ <- reportObservableSrv.create(ReportObservable(), job, richObservable.observable)
+                richObservable    <- observableSrv.create(artifact.toObservable(job._id, origObs.organisationIds), createdAttachment)
+                _                 <- reportObservableSrv.create(ReportObservable(), job, richObservable.observable)
               } yield createdAttachment
             }
           }
@@ -304,7 +294,7 @@ object JobOps {
     def can(permission: Permission)(implicit authContext: AuthContext): Traversal.V[Job] =
       if (authContext.permissions.contains(permission))
         traversal.filter(_.observable.can(permission))
-      else traversal.limit(0)
+      else traversal.empty
 
     def observable: Traversal.V[Observable] = traversal.in[ObservableJob].v[Observable]
 

migration/src/main/scala/org/thp/thehive/migration/Migrate.scala

@@ -205,6 +205,9 @@ object Migrate extends App with MigrationOps {
       migrationStats.flush()
       logger.info(migrationStats.toString)
       System.exit(returnStatus)
-    } finally actorSystem.terminate()
+    } finally {
+      actorSystem.terminate()
+      ()
+    }
   }
 }