From 7c1922f198ce7bd664e6870a6c9fa8649a5399b8 Mon Sep 17 00:00:00 2001 From: To-om Date: Tue, 9 Mar 2021 09:58:17 +0100 Subject: [PATCH] #1731 Fix tags in case template during migration --- ScalliGraph | 2 +- .../migration/dto/InputCaseTemplate.scala | 1 - .../thehive/migration/th3/Conversion.scala | 1 - .../thp/thehive/migration/th4/Output.scala | 58 +++++++++++++------ 4 files changed, 42 insertions(+), 20 deletions(-) diff --git a/ScalliGraph b/ScalliGraph index bad4e7e9e3..1cac59e0f9 160000 --- a/ScalliGraph +++ b/ScalliGraph @@ -1 +1 @@ -Subproject commit bad4e7e9e33e02a7cae93f7bbb5b9d940211cedc +Subproject commit 1cac59e0f9647e7907546d31c50d4ded172cdf7a diff --git a/migration/src/main/scala/org/thp/thehive/migration/dto/InputCaseTemplate.scala b/migration/src/main/scala/org/thp/thehive/migration/dto/InputCaseTemplate.scala index 76138e032c..684dd8950d 100644 --- a/migration/src/main/scala/org/thp/thehive/migration/dto/InputCaseTemplate.scala +++ b/migration/src/main/scala/org/thp/thehive/migration/dto/InputCaseTemplate.scala @@ -7,6 +7,5 @@ case class InputCaseTemplate( metaData: MetaData, caseTemplate: CaseTemplate, organisation: String, - tags: Set[String], customFields: Seq[InputCustomFieldValue] ) diff --git a/migration/src/main/scala/org/thp/thehive/migration/th3/Conversion.scala b/migration/src/main/scala/org/thp/thehive/migration/th3/Conversion.scala index 5992999f95..9adbc3f07b 100644 --- a/migration/src/main/scala/org/thp/thehive/migration/th3/Conversion.scala +++ b/migration/src/main/scala/org/thp/thehive/migration/th3/Conversion.scala @@ -390,7 +390,6 @@ trait Conversion { summary = summary ), mainOrganisation, - tags, (metricsValue ++ customFieldsValue).toSeq ) } diff --git a/migration/src/main/scala/org/thp/thehive/migration/th4/Output.scala b/migration/src/main/scala/org/thp/thehive/migration/th4/Output.scala index 105915ff96..46fd86af85 100644 --- a/migration/src/main/scala/org/thp/thehive/migration/th4/Output.scala +++ b/migration/src/main/scala/org/thp/thehive/migration/th4/Output.scala @@ -3,7 +3,7 @@ package org.thp.thehive.migration.th4 import akka.actor.ActorSystem import akka.actor.typed.Scheduler import akka.stream.Materializer -import com.google.inject.Guice +import com.google.inject.{Guice, Injector => GInjector} import net.codingwell.scalaguice.{ScalaModule, ScalaMultibinder} import org.apache.tinkerpop.gremlin.process.traversal.P import org.thp.scalligraph._ @@ -21,7 +21,6 @@ import org.thp.thehive.migration.dto._ import org.thp.thehive.models._ import org.thp.thehive.services._ import org.thp.thehive.{migration, ClusterSetup} -import play.api.cache.SyncCacheApi import play.api.cache.ehcache.EhCacheModule import play.api.inject.guice.GuiceInjector import play.api.inject.{ApplicationLifecycle, DefaultApplicationLifecycle, Injector} @@ -35,7 +34,7 @@ import scala.util.{Failure, Success, Try} object Output { - private def buildApp(configuration: Configuration)(implicit actorSystem: ActorSystem) = + private def buildApp(configuration: Configuration)(implicit actorSystem: ActorSystem): GInjector = Guice .createInjector( (play.api.inject.guice.GuiceableModule.guiceable(new EhCacheModule).guiced(Environment.simple(), configuration, Set.empty) :+ @@ -109,8 +108,7 @@ class Output @Inject() ( resolutionStatusSrv: ResolutionStatusSrv, jobSrv: JobSrv, actionSrv: ActionSrv, - db: Database, - cache: SyncCacheApi + db: Database ) extends migration.Output { lazy val logger: Logger = Logger(getClass) val defaultUserDomain: String = userSrv @@ -129,6 +127,7 @@ class Output @Inject() ( private var caseTemplates: Map[String, CaseTemplate with Entity] = Map.empty private var caseNumbers: Set[Int] = Set.empty private var alerts: Set[(String, String, String)] = Set.empty + private var tags: Map[String, Tag with Entity] = Map.empty private def retrieveExistingData(): Unit = { val profilesBuilder = Map.newBuilder[String, Profile with Entity] @@ -141,6 +140,7 @@ class Output @Inject() ( val caseTemplatesBuilder = Map.newBuilder[String, CaseTemplate with Entity] val caseNumbersBuilder = Set.newBuilder[Int] val alertsBuilder = Set.newBuilder[(String, String, String)] + val tagsBuilder = Map.newBuilder[String, Tag with Entity] db.roTransaction { implicit graph => graph @@ -157,7 +157,8 @@ class Output @Inject() ( "CustomField", "CaseTemplate", "Case", - "Alert" + "Alert", + "Tag" ) ) .toIterator @@ -194,6 +195,12 @@ class Output @Inject() ( val source = UMapping.string.getProperty(vertex, "source") val sourceRef = UMapping.string.getProperty(vertex, "sourceRef") alertsBuilder += ((`type`, source, sourceRef)) + case ("Tag", vertex) => + val tag = tagSrv.model.converter(vertex) + if (tag.namespace.startsWith(s"_freetags_")) + tagsBuilder += (s"${tag.namespace.drop(10)}-${tag.predicate}" -> tag) + else + tagsBuilder += (tag.toString -> tag) case _ => } } @@ -207,6 +214,7 @@ class Output @Inject() ( caseTemplates = caseTemplatesBuilder.result() caseNumbers = caseNumbersBuilder.result() alerts = alertsBuilder.result() + tags = tagsBuilder.result() if ( profiles.nonEmpty || organisations.nonEmpty || @@ -217,7 +225,8 @@ class Output @Inject() ( customFields.nonEmpty || caseTemplates.nonEmpty || caseNumbers.nonEmpty || - alerts.nonEmpty + alerts.nonEmpty || + tags.nonEmpty ) logger.info(s"""Already migrated: | ${profiles.size} profiles @@ -229,7 +238,8 @@ class Output @Inject() ( | ${customFields.size} customFields | ${caseTemplates.size} caseTemplates | ${caseNumbers.size} caseNumbers - | ${alerts.size} alerts""".stripMargin) + | ${alerts.size} alerts + | ${tags.size} tags""".stripMargin) } def startMigration(): Try[Unit] = @@ -282,9 +292,15 @@ class Output @Inject() ( } def getTag(tagName: String, organisationId: String)(implicit graph: Graph, authContext: AuthContext): Try[Tag with Entity] = - cache.getOrElseUpdate(s"tag-$tagName")( - tagSrv.createEntity(Tag(s"_freetags_$organisationId", tagName, None, None, tagSrv.freeTagColour)) - ) + tags + .get(tagName) + .orElse(tags.get(s"$organisationId-$tagName")) + .fold[Try[Tag with Entity]] { + tagSrv.createEntity(Tag(s"_freetags_$organisationId", tagName, None, None, tagSrv.freeTagColour)).map { tag => + tags += (tagName -> tag) + tag + } + }(Success.apply) override def organisationExists(inputOrganisation: InputOrganisation): Boolean = organisations.contains(inputOrganisation.organisation.name) @@ -462,19 +478,27 @@ class Output @Inject() ( authTransaction(inputCaseTemplate.metaData.createdBy) { implicit graph => implicit authContext => logger.debug(s"Create case template ${inputCaseTemplate.caseTemplate.name}") for { - organisation <- getOrganisation(inputCaseTemplate.organisation) - richCaseTemplate <- caseTemplateSrv.create(inputCaseTemplate.caseTemplate, organisation, Nil, Nil) - _ = updateMetaData(richCaseTemplate.caseTemplate, inputCaseTemplate.metaData) + organisation <- getOrganisation(inputCaseTemplate.organisation) + createdCaseTemplate <- caseTemplateSrv.createEntity(inputCaseTemplate.caseTemplate) + _ <- caseTemplateSrv.caseTemplateOrganisationSrv.create(CaseTemplateOrganisation(), createdCaseTemplate, organisation) + _ <- + inputCaseTemplate + .caseTemplate + .tags + .toTry( + getTag(_, organisation._id.value).flatMap(t => caseTemplateSrv.caseTemplateTagSrv.create(CaseTemplateTag(), createdCaseTemplate, t)) + ) + _ = updateMetaData(createdCaseTemplate, inputCaseTemplate.metaData) _ = inputCaseTemplate.customFields.foreach { case InputCustomFieldValue(name, value, order) => (for { cf <- getCustomField(name) ccf <- CustomFieldType.map(cf.`type`).setValue(CaseTemplateCustomField(order = order), value) - _ <- caseTemplateSrv.caseTemplateCustomFieldSrv.create(ccf, richCaseTemplate.caseTemplate, cf) + _ <- caseTemplateSrv.caseTemplateCustomFieldSrv.create(ccf, createdCaseTemplate, cf) } yield ()).logFailure(s"Unable to set custom field $name=${value.getOrElse("")}") } - _ = caseTemplates += (inputCaseTemplate.caseTemplate.name -> richCaseTemplate.caseTemplate) - } yield IdMapping(inputCaseTemplate.metaData.id, richCaseTemplate._id) + _ = caseTemplates += (inputCaseTemplate.caseTemplate.name -> createdCaseTemplate) + } yield IdMapping(inputCaseTemplate.metaData.id, createdCaseTemplate._id) } override def createCaseTemplateTask(caseTemplateId: EntityId, inputTask: InputTask): Try[IdMapping] =