#1501 Add type to safely handle entity ID

cortex/connector/src/main/scala/org/thp/thehive/connector/cortex/controllers/v0/ActionCtrl.scala

@@ -1,6 +1,7 @@
 package org.thp.thehive.connector.cortex.controllers.v0
 
 import javax.inject.{Inject, Named, Singleton}
+import org.thp.scalligraph.EntityIdOrName
 import org.thp.scalligraph.auth.AuthContext
 import org.thp.scalligraph.controllers.{Entrypoint, FieldsParser}
 import org.thp.scalligraph.models.{Database, Entity, UMapping}
@@ -57,7 +58,7 @@ class ActionCtrl @Inject() (
       .asyncAuth { implicit request =>
         val action: InputAction = request.body("action")
         val tryEntity = db.roTransaction { implicit graph =>
-          entityHelper.get(toObjectType(action.objectType), action.objectId, Permissions.manageAction)
+          entityHelper.get(toObjectType(action.objectType), EntityIdOrName(action.objectId), Permissions.manageAction)
         }
         for {
           entity <- Future.fromTry(tryEntity)
@@ -69,21 +70,21 @@ class ActionCtrl @Inject() (
     entrypoint("get by entity")
       .authRoTransaction(db) { implicit request => implicit graph =>
         for {
-          entity <- entityHelper.get(toObjectType(objectType), objectId, Permissions.manageAction)
+          entity <- entityHelper.get(toObjectType(objectType), EntityIdOrName(objectId), Permissions.manageAction)
         } yield Results.Ok(actionSrv.listForEntity(entity._id).toJson)
       }
 }
 
 @Singleton
-class PublicAction @Inject() (actionSrv: ActionSrv) extends PublicData {
+class PublicAction @Inject() (actionSrv: ActionSrv, @Named("with-thehive-schema") db: Database) extends PublicData {
 
   override val entityName: String = "action"
   override val initialQuery: Query =
     Query.init[Traversal.V[Action]]("listAction", (graph, authContext) => actionSrv.startTraversal(graph).visible(authContext))
-  override val getQuery: ParamQuery[IdOrName] = Query.initWithParam[IdOrName, Traversal.V[Action]](
+  override val getQuery: ParamQuery[EntityIdOrName] = Query.initWithParam[EntityIdOrName, Traversal.V[Action]](
     "getAction",
-    FieldsParser[IdOrName],
-    (param, graph, authContext) => actionSrv.get(param.idOrName)(graph).visible(authContext)
+    FieldsParser[EntityIdOrName],
+    (idOrName, graph, authContext) => actionSrv.get(idOrName)(graph).visible(authContext)
   )
   override val pageQuery: ParamQuery[OutputParam] = Query.withParam[OutputParam, Traversal.V[Action], IteratorOutput](
     "page",
@@ -111,7 +112,7 @@ class PublicAction @Inject() (actionSrv: ActionSrv) extends PublicData {
       .property("objectType", UMapping.string)(_.select(_.context.domainMap(o => fromObjectType(o._label))).readonly)
       .property("status", UMapping.string)(_.field.readonly)
       .property("startDate", UMapping.date)(_.field.readonly)
-      .property("objectId", UMapping.id)(_.select(_.out[ActionContext]._id).readonly)
+      .property("objectId", db.idMapping)(_.select(_.out[ActionContext]._id).readonly)
       .property("responderName", UMapping.string.optional)(_.field.readonly)
       .property("cortexId", UMapping.string.optional)(_.field.readonly)
       .property("tlp", UMapping.int.optional)(_.field.readonly)

cortex/connector/src/main/scala/org/thp/thehive/connector/cortex/controllers/v0/AnalyzerTemplateCtrl.scala

@@ -4,6 +4,7 @@ import java.util.zip.ZipFile
 
 import com.google.inject.name.Named
 import javax.inject.{Inject, Singleton}
+import org.thp.scalligraph.EntityIdOrName
 import org.thp.scalligraph.controllers.{Entrypoint, FFile, FieldsParser}
 import org.thp.scalligraph.models.{Database, Entity, UMapping}
 import org.thp.scalligraph.query._
@@ -12,9 +13,10 @@ import org.thp.scalligraph.traversal.{IteratorOutput, Traversal}
 import org.thp.thehive.connector.cortex.controllers.v0.Conversion._
 import org.thp.thehive.connector.cortex.dto.v0.InputAnalyzerTemplate
 import org.thp.thehive.connector.cortex.models.AnalyzerTemplate
+import org.thp.thehive.connector.cortex.services.AnalyzerTemplateOps._
 import org.thp.thehive.connector.cortex.services.AnalyzerTemplateSrv
 import org.thp.thehive.controllers.v0.Conversion._
-import org.thp.thehive.controllers.v0.{IdOrName, OutputParam, PublicData, QueryCtrl}
+import org.thp.thehive.controllers.v0.{OutputParam, PublicData, QueryCtrl}
 import org.thp.thehive.models.Permissions
 import play.api.libs.json.{JsFalse, JsObject, JsTrue}
 import play.api.mvc.{Action, AnyContent, Results}
@@ -34,7 +36,7 @@ class AnalyzerTemplateCtrl @Inject() (
     entrypoint("get content")
       .authTransaction(db) { _ => implicit graph =>
         analyzerTemplateSrv
-          .getOrFail(id)
+          .getOrFail(EntityIdOrName(id))
           .map(report => Results.Ok(report.content))
       }
 
@@ -69,7 +71,7 @@ class AnalyzerTemplateCtrl @Inject() (
     entrypoint("delete template")
       .authPermittedTransaction(db, Permissions.manageAnalyzerTemplate) { implicit request => implicit graph =>
         analyzerTemplateSrv
-          .get(id)
+          .get(EntityIdOrName(id))
           .getOrFail("AnalyzerTemplate")
           .map { analyzerTemplate =>
             analyzerTemplateSrv.remove(analyzerTemplate)
@@ -84,7 +86,7 @@ class AnalyzerTemplateCtrl @Inject() (
         val propertyUpdaters: Seq[PropertyUpdater] = request.body("template")
 
         for {
-          (templateSteps, _) <- analyzerTemplateSrv.update(_.getByIds(id), propertyUpdaters)
+          (templateSteps, _) <- analyzerTemplateSrv.update(_.get(EntityIdOrName(id)), propertyUpdaters)
           template           <- templateSteps.getOrFail("AnalyzerTemplate")
         } yield Results.Ok(template.toJson)
 
@@ -96,10 +98,10 @@ class PublicAnalyzerTemplate @Inject() (analyzerTemplateSrv: AnalyzerTemplateSrv
   override val entityName: String = "analyzerTemplate"
   override val initialQuery: Query =
     Query.init[Traversal.V[AnalyzerTemplate]]("listAnalyzerTemplate", (graph, _) => analyzerTemplateSrv.startTraversal(graph))
-  override val getQuery: ParamQuery[IdOrName] = Query.initWithParam[IdOrName, Traversal.V[AnalyzerTemplate]](
+  override val getQuery: ParamQuery[EntityIdOrName] = Query.initWithParam[EntityIdOrName, Traversal.V[AnalyzerTemplate]](
     "getReportTemplate",
-    FieldsParser[IdOrName],
-    (param, graph, _) => analyzerTemplateSrv.get(param.idOrName)(graph)
+    FieldsParser[EntityIdOrName],
+    (idOrName, graph, _) => analyzerTemplateSrv.get(idOrName)(graph)
   )
   override val pageQuery: ParamQuery[OutputParam] =
     Query.withParam[OutputParam, Traversal.V[AnalyzerTemplate], IteratorOutput](

cortex/connector/src/main/scala/org/thp/thehive/connector/cortex/controllers/v0/Conversion.scala

@@ -18,7 +18,7 @@ object Conversion {
       .withFieldRenamed(_.workerName, _.responderName)
       .withFieldRenamed(_.workerDefinition, _.responderDefinition)
       .withFieldComputed(_.status, _.status.toString)
-      .withFieldComputed(_.objectId, _.context._id)
+      .withFieldComputed(_.objectId, _.context._id.toString)
       .withFieldComputed(_.objectType, _.context._label)
       .withFieldComputed(_.operations, a => JsArray(a.operations).toString)
       .withFieldComputed(_.report, _.report.map(_.toString).getOrElse("{}"))
@@ -43,7 +43,7 @@ object Conversion {
             case r                                  => r + ("success" -> JsFalse)
           }
       )
-      .withFieldRenamed(_._id, _.id)
+      .withFieldComputed(_.id, _._id.toString)
       .withFieldConst(_._type, "case_artifact_job")
       .withFieldConst(_.case_artifact, None)
       .transform
@@ -69,7 +69,7 @@ object Conversion {
               case r                                  => r + ("success" -> JsFalse)
             }
         )
-        .withFieldRenamed(_._id, _.id)
+        .withFieldComputed(_.id, _._id.toString)
         .withFieldConst(_._type, "case_artifact_job")
         .withFieldConst(
           _.case_artifact,
@@ -86,7 +86,7 @@ object Conversion {
       at.asInstanceOf[AnalyzerTemplate]
         .into[OutputAnalyzerTemplate]
         .withFieldComputed(_.analyzerId, _.workerId)
-        .withFieldConst(_.id, at._id)
+        .withFieldConst(_.id, at._id.toString)
         .withFieldComputed(_.content, _.content)
         .transform
     )

cortex/connector/src/main/scala/org/thp/thehive/connector/cortex/controllers/v0/CortexQueryExecutor.scala

@@ -3,17 +3,18 @@ package org.thp.thehive.connector.cortex.controllers.v0
 import com.google.inject.name.Named
 import javax.inject.{Inject, Singleton}
 import org.scalactic.Good
-import org.thp.scalligraph.BadRequestError
 import org.thp.scalligraph.auth.AuthContext
 import org.thp.scalligraph.controllers.FieldsParser
 import org.thp.scalligraph.models._
 import org.thp.scalligraph.query._
 import org.thp.scalligraph.traversal.Traversal
 import org.thp.scalligraph.traversal.TraversalOps._
+import org.thp.scalligraph.{BadRequestError, EntityIdOrName}
 import org.thp.thehive.connector.cortex.models.Job
 import org.thp.thehive.connector.cortex.services.JobOps._
 import org.thp.thehive.controllers.v0._
 import org.thp.thehive.models.Observable
+import org.thp.thehive.services.ObservableOps._
 
 import scala.reflect.runtime.{universe => ru}
 
@@ -66,7 +67,7 @@ class CortexParentIdInputFilter(parentId: String) extends InputQuery[Traversal.U
       authContext: AuthContext
   ): Traversal.Unk =
     if (traversalType =:= ru.typeOf[Traversal.V[Job]])
-      traversal.asInstanceOf[Traversal.V[Job]].filter(_.observable.getByIds(parentId)).asInstanceOf[Traversal.Unk]
+      traversal.asInstanceOf[Traversal.V[Job]].filter(_.observable.get(EntityIdOrName(parentId))).asInstanceOf[Traversal.Unk]
     else throw BadRequestError(s"$traversalType hasn't parent")
 }
 

cortex/connector/src/main/scala/org/thp/thehive/connector/cortex/controllers/v0/JobCtrl.scala

@@ -7,13 +7,13 @@ import org.thp.scalligraph.models.{Database, UMapping}
 import org.thp.scalligraph.query._
 import org.thp.scalligraph.traversal.TraversalOps._
 import org.thp.scalligraph.traversal.{IteratorOutput, Traversal}
-import org.thp.scalligraph.{AuthorizationError, ErrorHandler}
+import org.thp.scalligraph.{AuthorizationError, EntityIdOrName, ErrorHandler}
 import org.thp.thehive.connector.cortex.controllers.v0.Conversion._
 import org.thp.thehive.connector.cortex.models.{Job, RichJob}
 import org.thp.thehive.connector.cortex.services.JobOps._
 import org.thp.thehive.connector.cortex.services.JobSrv
 import org.thp.thehive.controllers.v0.Conversion._
-import org.thp.thehive.controllers.v0.{IdOrName, OutputParam, PublicData, QueryCtrl}
+import org.thp.thehive.controllers.v0.{OutputParam, PublicData, QueryCtrl}
 import org.thp.thehive.models.{Permissions, RichCase, RichObservable}
 import org.thp.thehive.services.ObservableOps._
 import org.thp.thehive.services.ObservableSrv
@@ -36,7 +36,7 @@ class JobCtrl @Inject() (
     entrypoint("get job")
       .authRoTransaction(db) { implicit request => implicit graph =>
         jobSrv
-          .getByIds(jobId)
+          .get(EntityIdOrName(jobId))
           .visible
           .richJob
           .getOrFail("Job")
@@ -55,8 +55,8 @@ class JobCtrl @Inject() (
           db.roTransaction { implicit graph =>
             val artifactId: String = request.body("artifactId")
             for {
-              o <- observableSrv.getByIds(artifactId).richObservable.getOrFail("Observable")
-              c <- observableSrv.getByIds(artifactId).`case`.getOrFail("Case")
+              o <- observableSrv.get(EntityIdOrName(artifactId)).richObservable.getOrFail("Observable")
+              c <- observableSrv.get(EntityIdOrName(artifactId)).`case`.getOrFail("Case")
             } yield (o, c)
           }.fold(
             error => errorHandler.onServerError(request, error),
@@ -76,10 +76,10 @@ class PublicJob @Inject() (jobSrv: JobSrv) extends PublicData with JobRenderer {
   override val entityName: String = "job"
   override val initialQuery: Query =
     Query.init[Traversal.V[Job]]("listJob", (graph, authContext) => jobSrv.startTraversal(graph).visible(authContext))
-  override val getQuery: ParamQuery[IdOrName] = Query.initWithParam[IdOrName, Traversal.V[Job]](
+  override val getQuery: ParamQuery[EntityIdOrName] = Query.initWithParam[EntityIdOrName, Traversal.V[Job]](
     "getJob",
-    FieldsParser[IdOrName],
-    (param, graph, authContext) => jobSrv.get(param.idOrName)(graph).visible(authContext)
+    FieldsParser[EntityIdOrName],
+    (idOrName, graph, authContext) => jobSrv.get(idOrName)(graph).visible(authContext)
   )
   override val pageQuery: ParamQuery[OutputParam] =
     Query.withParam[OutputParam, Traversal.V[Job], IteratorOutput](

cortex/connector/src/main/scala/org/thp/thehive/connector/cortex/controllers/v0/ResponderCtrl.scala

@@ -2,6 +2,7 @@ package org.thp.thehive.connector.cortex.controllers.v0
 
 import com.google.inject.name.Named
 import javax.inject.{Inject, Singleton}
+import org.thp.scalligraph.EntityIdOrName
 import org.thp.scalligraph.controllers.{Entrypoint, FieldsParser}
 import org.thp.scalligraph.models.Database
 import org.thp.thehive.connector.cortex.controllers.v0.Conversion._
@@ -24,7 +25,7 @@ class ResponderCtrl @Inject() (
     entrypoint("get responders")
       .asyncAuth { implicit req =>
         responderSrv
-          .getRespondersByType(entityType, entityId)
+          .getRespondersByType(entityType, EntityIdOrName(entityId))
           .map(l => Results.Ok(l.toSeq.toJson))
       }
 

cortex/connector/src/main/scala/org/thp/thehive/connector/cortex/controllers/v0/Router.scala

@@ -10,7 +10,6 @@ class Router @Inject() (
     val jobCtrl: JobCtrl,
     analyzerCtrl: AnalyzerCtrl,
     val actionCtrl: ActionCtrl,
-    cortexQueryExecutor: CortexQueryExecutor,
     val reportCtrl: AnalyzerTemplateCtrl,
     responderCtrl: ResponderCtrl
 ) extends SimpleRouter {

cortex/connector/src/main/scala/org/thp/thehive/connector/cortex/models/Action.scala

@@ -3,9 +3,9 @@ package org.thp.thehive.connector.cortex.models
 import java.util.Date
 
 import org.apache.tinkerpop.gremlin.structure.{Edge, Graph, Vertex}
-import org.thp.scalligraph.BuildVertexEntity
 import org.thp.scalligraph.models._
 import org.thp.scalligraph.traversal.Converter
+import org.thp.scalligraph.{BuildVertexEntity, EntityId}
 import play.api.libs.json.JsObject
 
 @BuildVertexEntity
@@ -24,7 +24,7 @@ case class Action(
 )
 
 case class RichAction(action: Action with Entity, context: Product with Entity) {
-  def _id: String               = action._id
+  def _id: EntityId             = action._id
   def _createdAt: Date          = action._createdAt
   def _createdBy: String        = action._createdBy
   def workerId: String          = action.workerId
@@ -51,7 +51,7 @@ object ActionContext extends HasModel {
     override val fields: Map[String, Mapping[_, _, _]] = Map.empty
     override val converter: Converter[EEntity, Edge] = (element: Edge) =>
       new ActionContext with Entity {
-        override val _id: String                = element.id().toString
+        override val _id: EntityId              = EntityId(element.id())
         override val _label: String             = "ActionContext"
         override val _createdBy: String         = UMapping.string.getProperty(element, "_createdBy")
         override val _updatedBy: Option[String] = UMapping.string.optional.getProperty(element, "_updatedBy")
@@ -60,7 +60,7 @@ object ActionContext extends HasModel {
       }
     override def addEntity(a: ActionContext, entity: Entity): EEntity =
       new ActionContext with Entity {
-        override def _id: String                = entity._id
+        override def _id: EntityId              = entity._id
         override def _label: String             = entity._label
         override def _createdBy: String         = entity._createdBy
         override def _updatedBy: Option[String] = entity._updatedBy

cortex/connector/src/main/scala/org/thp/thehive/connector/cortex/models/Job.scala

@@ -3,7 +3,7 @@ package org.thp.thehive.connector.cortex.models
 import java.util.Date
 
 import org.thp.scalligraph.models.Entity
-import org.thp.scalligraph.{BuildEdgeEntity, BuildVertexEntity}
+import org.thp.scalligraph.{BuildEdgeEntity, BuildVertexEntity, EntityId}
 import org.thp.thehive.models.{Observable, RichObservable}
 import play.api.libs.json.{Format, JsObject, Json}
 
@@ -36,7 +36,7 @@ case class RichJob(
     job: Job with Entity,
     observables: Seq[(RichObservable, JsObject)]
 ) {
-  def _id: String                = job._id
+  def _id: EntityId              = job._id
   def _createdBy: String         = job._createdBy
   def _updatedBy: Option[String] = job._updatedBy
   def _createdAt: Date           = job._createdAt

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

@@ -4,10 +4,10 @@ import java.util.Date
 
 import javax.inject.Inject
 import org.apache.tinkerpop.gremlin.structure.Graph
-import org.thp.scalligraph.InternalError
 import org.thp.scalligraph.auth.AuthContext
 import org.thp.scalligraph.models.Entity
 import org.thp.scalligraph.traversal.TraversalOps._
+import org.thp.scalligraph.{EntityIdOrName, InternalError}
 import org.thp.thehive.connector.cortex.models._
 import org.thp.thehive.controllers.v0.Conversion._
 import org.thp.thehive.dto.v0.InputTask
@@ -41,8 +41,8 @@ class ActionOperationSrv @Inject() (
     * @param authContext auth for access check
     * @return
     */
-  def execute(entity: Entity, operation: ActionOperation, relatedCase: Option[Case with Entity], relatedTask: Option[Task with Entity])(
-      implicit graph: Graph,
+  def execute(entity: Entity, operation: ActionOperation, relatedCase: Option[Case with Entity], relatedTask: Option[Task with Entity])(implicit
+      graph: Graph,
       authContext: AuthContext
   ): Try[ActionOperationStatus] = {
 
@@ -72,7 +72,7 @@ class ActionOperationSrv @Inject() (
       case AddCustomFields(name, _, value) =>
         for {
           c <- relatedCase.fold[Try[Case with Entity]](Failure(InternalError("Unable to apply action AddCustomFields without case")))(Success(_))
-          _ <- caseSrv.setOrCreateCustomField(c, name, Some(value), None)
+          _ <- caseSrv.setOrCreateCustomField(c, EntityIdOrName(name), Some(value), None)
         } yield updateOperation(operation)
 
       case CloseTask() =>
@@ -96,7 +96,7 @@ class ActionOperationSrv @Inject() (
       case AddArtifactToCase(_, dataType, dataMessage) =>
         for {
           c       <- relatedCase.fold[Try[Case with Entity]](Failure(InternalError("Unable to apply action AddArtifactToCase without case")))(Success(_))
-          obsType <- observableTypeSrv.getOrFail(dataType)
+          obsType <- observableTypeSrv.getOrFail(EntityIdOrName(dataType))
           richObservable <- observableSrv.create(
             Observable(Some(dataMessage), 2, ioc = false, sighted = false),
             obsType,
@@ -110,8 +110,8 @@ class ActionOperationSrv @Inject() (
       case AssignCase(owner) =>
         for {
           c <- relatedCase.fold[Try[Case with Entity]](Failure(InternalError("Unable to apply action AssignCase without case")))(Success(_))
-          u <- userSrv.get(owner).getOrFail("User")
-          _ <- Try(caseSrv.startTraversal.getByIds(c._id).unassign())
+          u <- userSrv.get(EntityIdOrName(owner)).getOrFail("User")
+          _ <- Try(caseSrv.startTraversal.getEntity(c).unassign())
           _ <- caseSrv.assign(c, u)
         } yield updateOperation(operation)