#1731 Update migration tool for TheHive 3.5

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

Hide file tree

Showing 14 changed files with 323 additions and 303 deletions.
diff --git a/ScalliGraph b/ScalliGraph
diff --git a/build.sbt b/build.sbt
@@ -62,7 +62,7 @@ libraryDependencies in ThisBuild ++= {
 }
 dependencyOverrides in ThisBuild ++= Seq(
 //  "org.locationtech.spatial4j" % "spatial4j"                 % "0.6",
-  "org.elasticsearch.client" % "elasticsearch-rest-client" % "6.7.2"
+//  "org.elasticsearch.client" % "elasticsearch-rest-client" % "6.7.2"
 )
 PlayKeys.includeDocumentationInBinary := false
 milestoneFilter := ((milestone: Milestone) => milestone.title.startsWith("4"))
@@ -337,12 +337,13 @@ lazy val thehiveMigration = (project in file("migration"))
     libraryDependencies ++= Seq(
       elastic4sCore,
       elastic4sHttpStreams,
-      elastic4sHttp,
+      elastic4sClient,
 //      jts,
       ehcache,
       scopt,
       specs % Test
     ),
+    dependencyOverrides += akkaActor,
     fork := true,
     normalizedName := "migrate"
   )

diff --git a/migration/src/main/scala/org/thp/thehive/migration/dto/InputAlert.scala b/migration/src/main/scala/org/thp/thehive/migration/dto/InputAlert.scala
@@ -7,7 +7,6 @@ case class InputAlert(
     alert: Alert,
     caseId: Option[String],
     organisation: String,
-    tags: Set[String],
     customFields: Map[String, Option[Any]],
     caseTemplate: Option[String]
 ) {

diff --git a/migration/src/main/scala/org/thp/thehive/migration/dto/InputCase.scala b/migration/src/main/scala/org/thp/thehive/migration/dto/InputCase.scala
@@ -4,12 +4,7 @@ import org.thp.thehive.models.Case
 
 case class InputCase(
     `case`: Case,
-    user: Option[String],
     organisations: Map[String, String],
-    tags: Set[String],
     customFields: Map[String, Option[Any]],
-    caseTemplate: Option[String],
-    resolutionStatus: Option[String],
-    impactStatus: Option[String],
     metaData: MetaData
 )
diff --git a/migration/src/main/scala/org/thp/thehive/migration/dto/InputObservable.scala b/migration/src/main/scala/org/thp/thehive/migration/dto/InputObservable.scala
@@ -6,7 +6,5 @@ case class InputObservable(
     metaData: MetaData,
     observable: Observable,
     organisations: Seq[String],
-    `type`: String,
-    tags: Set[String],
     dataOrAttachment: Either[String, InputAttachment]
 )
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
@@ -1,11 +1,8 @@
 package org.thp.thehive.migration.th3
 
-import java.util.{Base64, Date}
-
 import akka.NotUsed
 import akka.stream.scaladsl.Source
 import akka.util.ByteString
-import org.thp.scalligraph.EntityId
 import org.thp.scalligraph.utils.Hash
 import org.thp.thehive.connector.cortex.models.{Action, Job, JobStatus}
 import org.thp.thehive.controllers.v0
@@ -15,6 +12,8 @@ import org.thp.thehive.models._
 import play.api.libs.functional.syntax._
 import play.api.libs.json._
 
+import java.util.{Base64, Date}
+
 case class Attachment(name: String, hashes: Seq[Hash], size: Long, contentType: String, id: String)
 trait Conversion {
 
@@ -64,7 +63,7 @@ trait Conversion {
       status      <- (json \ "status").validate[CaseStatus.Value]
       summary     <- (json \ "summary").validateOpt[String]
       user        <- (json \ "owner").validateOpt[String]
-      tags             = (json \ "tags").asOpt[Set[String]].getOrElse(Set.empty)
+      tags             = (json \ "tags").asOpt[Set[String]].getOrElse(Set.empty).filterNot(_.isEmpty)
       metrics          = (json \ "metrics").asOpt[JsObject].getOrElse(JsObject.empty)
       resolutionStatus = (json \ "resolutionStatus").asOpt[String]
       impactStatus     = (json \ "impactStatus").asOpt[String]
@@ -91,18 +90,13 @@ trait Conversion {
         tags = tags.toSeq,
         number = number,
         organisationIds = Nil,
-        assignee = None,
+        assignee = user.map(normaliseLogin),
         impactStatus = impactStatus,
         resolutionStatus = resolutionStatus,
         caseTemplate = None
       ), // organisation Ids are filled by output
-      user.map(normaliseLogin),
       Map(mainOrganisation -> Profile.orgAdmin.name),
-      tags,
       (metricsValue ++ customFieldsValue).toMap,
-      None,
-      resolutionStatus,
-      impactStatus,
       metaData
     )
   }
@@ -133,13 +127,10 @@ trait Conversion {
         ioc = ioc,
         sighted = sighted,
         ignoreSimilarity = None,
-        data = dataOrAttachment.swap.toOption,
         dataType = dataType,
         tags = tags.toSeq
       ),
       Seq(mainOrganisation),
-      dataType,
-      tags,
       dataOrAttachment
     )
   }
@@ -207,7 +198,7 @@ trait Conversion {
       read = status == "Ignored" || status == "Imported"
       follow <- (json \ "follow").validate[Boolean]
       caseId <- (json \ "case").validateOpt[String]
-      tags         = (json \ "tags").asOpt[Set[String]].getOrElse(Set.empty)
+      tags         = (json \ "tags").asOpt[Set[String]].getOrElse(Set.empty).filterNot(_.isEmpty)
       customFields = (json \ "metrics").asOpt[JsObject].getOrElse(JsObject.empty)
       customFieldsValue = customFields.value.map {
         case (name, value) =>
@@ -234,7 +225,6 @@ trait Conversion {
       ),
       caseId,
       mainOrganisation,
-      tags,
       customFieldsValue.toMap,
       caseTemplate: Option[String]
     )
@@ -265,15 +255,10 @@ trait Conversion {
           ioc = ioc.getOrElse(false),
           sighted = false,
           ignoreSimilarity = None,
-          data = dataOrAttachment.swap.toOption,
           dataType = dataType,
-          tags = tags.toSeq,
-          organisationIds = Nil,
-          relatedId = EntityId("")
+          tags = tags.toSeq
         ),
-        Nil,
-        dataType,
-        tags,
+        Seq(mainOrganisation),
         dataOrAttachment
       )
 
@@ -497,13 +482,10 @@ trait Conversion {
           ioc = ioc,
           sighted = sighted,
           ignoreSimilarity = None,
-          data = dataOrAttachment.swap.toOption,
           dataType = dataType,
           tags = tags.toSeq
         ),
         Seq(mainOrganisation),
-        dataType,
-        tags,
         dataOrAttachment
       )
     }

diff --git a/migration/src/main/scala/org/thp/thehive/migration/th3/DBConfiguration.scala b/migration/src/main/scala/org/thp/thehive/migration/th3/DBConfiguration.scala
@@ -2,16 +2,16 @@ package org.thp.thehive.migration.th3
 
 import java.nio.file.{Files, Paths}
 import java.security.KeyStore
-
 import akka.NotUsed
 import akka.actor.ActorSystem
 import akka.stream.scaladsl.{Sink, Source}
-import com.sksamuel.elastic4s.http._
-import com.sksamuel.elastic4s.http.bulk.BulkResponseItem
-import com.sksamuel.elastic4s.http.search.SearchHit
-import com.sksamuel.elastic4s.searches._
+import com.sksamuel.elastic4s._
+import com.sksamuel.elastic4s.http.JavaClient
+import com.sksamuel.elastic4s.requests.bulk.BulkResponseItem
+import com.sksamuel.elastic4s.requests.searches.{SearchHit, SearchRequest}
 import com.sksamuel.elastic4s.streams.ReactiveElastic.ReactiveElastic
 import com.sksamuel.elastic4s.streams.{RequestBuilder, ResponseListener}
+
 import javax.inject.{Inject, Named, Singleton}
 import javax.net.ssl.{KeyManagerFactory, SSLContext, TrustManagerFactory}
 import org.apache.http.auth.{AuthScope, UsernamePasswordCredentials}
@@ -39,29 +39,29 @@ class DBConfiguration @Inject() (
     config: Configuration,
     lifecycle: ApplicationLifecycle,
     @Named("databaseVersion") val version: Int,
-    implicit val ec: ExecutionContext,
     implicit val actorSystem: ActorSystem
 ) {
   private[DBConfiguration] lazy val logger = Logger(getClass)
 
-  def requestConfigCallback: RequestConfigCallback = (requestConfigBuilder: RequestConfig.Builder) => {
-    requestConfigBuilder.setAuthenticationEnabled(credentialsProviderMaybe.isDefined)
-    config.getOptional[Boolean]("search.circularRedirectsAllowed").foreach(requestConfigBuilder.setCircularRedirectsAllowed)
-    config.getOptional[Int]("search.connectionRequestTimeout").foreach(requestConfigBuilder.setConnectionRequestTimeout)
-    config.getOptional[Int]("search.connectTimeout").foreach(requestConfigBuilder.setConnectTimeout)
-    config.getOptional[Boolean]("search.contentCompressionEnabled").foreach(requestConfigBuilder.setContentCompressionEnabled)
-    config.getOptional[String]("search.cookieSpec").foreach(requestConfigBuilder.setCookieSpec)
-    config.getOptional[Boolean]("search.expectContinueEnabled").foreach(requestConfigBuilder.setExpectContinueEnabled)
-    //    config.getOptional[InetAddress]("search.localAddress").foreach(requestConfigBuilder.setLocalAddress)
-    config.getOptional[Int]("search.maxRedirects").foreach(requestConfigBuilder.setMaxRedirects)
-    //    config.getOptional[Boolean]("search.proxy").foreach(requestConfigBuilder.setProxy)
-    config.getOptional[Seq[String]]("search.proxyPreferredAuthSchemes").foreach(v => requestConfigBuilder.setProxyPreferredAuthSchemes(v.asJava))
-    config.getOptional[Boolean]("search.redirectsEnabled").foreach(requestConfigBuilder.setRedirectsEnabled)
-    config.getOptional[Boolean]("search.relativeRedirectsAllowed").foreach(requestConfigBuilder.setRelativeRedirectsAllowed)
-    config.getOptional[Int]("search.socketTimeout").foreach(requestConfigBuilder.setSocketTimeout)
-    config.getOptional[Seq[String]]("search.targetPreferredAuthSchemes").foreach(v => requestConfigBuilder.setTargetPreferredAuthSchemes(v.asJava))
-    requestConfigBuilder
-  }
+  def requestConfigCallback: RequestConfigCallback =
+    (requestConfigBuilder: RequestConfig.Builder) => {
+      requestConfigBuilder.setAuthenticationEnabled(credentialsProviderMaybe.isDefined)
+      config.getOptional[Boolean]("search.circularRedirectsAllowed").foreach(requestConfigBuilder.setCircularRedirectsAllowed)
+      config.getOptional[Int]("search.connectionRequestTimeout").foreach(requestConfigBuilder.setConnectionRequestTimeout)
+      config.getOptional[Int]("search.connectTimeout").foreach(requestConfigBuilder.setConnectTimeout)
+      config.getOptional[Boolean]("search.contentCompressionEnabled").foreach(requestConfigBuilder.setContentCompressionEnabled)
+      config.getOptional[String]("search.cookieSpec").foreach(requestConfigBuilder.setCookieSpec)
+      config.getOptional[Boolean]("search.expectContinueEnabled").foreach(requestConfigBuilder.setExpectContinueEnabled)
+      //    config.getOptional[InetAddress]("search.localAddress").foreach(requestConfigBuilder.setLocalAddress)
+      config.getOptional[Int]("search.maxRedirects").foreach(requestConfigBuilder.setMaxRedirects)
+      //    config.getOptional[Boolean]("search.proxy").foreach(requestConfigBuilder.setProxy)
+      config.getOptional[Seq[String]]("search.proxyPreferredAuthSchemes").foreach(v => requestConfigBuilder.setProxyPreferredAuthSchemes(v.asJava))
+      config.getOptional[Boolean]("search.redirectsEnabled").foreach(requestConfigBuilder.setRedirectsEnabled)
+      config.getOptional[Boolean]("search.relativeRedirectsAllowed").foreach(requestConfigBuilder.setRelativeRedirectsAllowed)
+      config.getOptional[Int]("search.socketTimeout").foreach(requestConfigBuilder.setSocketTimeout)
+      config.getOptional[Seq[String]]("search.targetPreferredAuthSchemes").foreach(v => requestConfigBuilder.setTargetPreferredAuthSchemes(v.asJava))
+      requestConfigBuilder
+    }
 
   lazy val credentialsProviderMaybe: Option[CredentialsProvider] =
     for {
@@ -86,9 +86,7 @@ class DBConfiguration @Inject() (
         val kmf = KeyManagerFactory.getInstance(KeyManagerFactory.getDefaultAlgorithm)
         kmf.init(keyStore, keyStorePassword)
         kmf.getKeyManagers
-      } finally {
-        keyInputStream.close()
-      }
+      } finally keyInputStream.close()
 
     val trustManagers = config
       .getOptional[String]("search.trustStore.path")
@@ -102,9 +100,7 @@ class DBConfiguration @Inject() (
           val tmf = TrustManagerFactory.getInstance(TrustManagerFactory.getDefaultAlgorithm)
           tmf.init(keyStore)
           tmf.getTrustManagers
-        } finally {
-          trustInputStream.close()
-        }
+        } finally trustInputStream.close()
       }
       .getOrElse(Array.empty)
 
@@ -114,27 +110,29 @@ class DBConfiguration @Inject() (
     sslContext
   }
 
-  def httpClientConfig: HttpClientConfigCallback = (httpClientBuilder: HttpAsyncClientBuilder) => {
-    sslContextMaybe.foreach(httpClientBuilder.setSSLContext)
-    credentialsProviderMaybe.foreach(httpClientBuilder.setDefaultCredentialsProvider)
-    httpClientBuilder
-  }
+  def httpClientConfig: HttpClientConfigCallback =
+    (httpClientBuilder: HttpAsyncClientBuilder) => {
+      sslContextMaybe.foreach(httpClientBuilder.setSSLContext)
+      credentialsProviderMaybe.foreach(httpClientBuilder.setDefaultCredentialsProvider)
+      httpClientBuilder
+    }
 
   /**
     * Underlying ElasticSearch client
     */
-  val client: ElasticClient = ElasticClient(ElasticProperties(config.get[String]("search.uri")), requestConfigCallback, httpClientConfig)
+  private val props  = ElasticProperties(config.get[String]("search.uri"))
+  private val client = ElasticClient(JavaClient(props, requestConfigCallback, httpClientConfig))
+
   // when application close, close also ElasticSearch connection
   lifecycle.addStopHook { () =>
-    Future {
-      client.close()
-    }
+    client.close()
+    Future.successful(())
   }
 
-  def execute[T, U](t: T)(
-      implicit
+  def execute[T, U](t: T)(implicit
       handler: Handler[T, U],
-      manifest: Manifest[U]
+      manifest: Manifest[U],
+      ec: ExecutionContext
   ): Future[U] = {
     logger.debug(s"Elasticsearch request: ${client.show(t)}")
     client.execute(t).flatMap {
@@ -157,12 +155,13 @@ class DBConfiguration @Inject() (
   /**
     * Creates a Source (akka stream) from the result of the search
     */
-  def source(searchRequest: SearchRequest): Source[SearchHit, NotUsed] = Source.fromPublisher(client.publisher(searchRequest))
+  def source(searchRequest: SearchRequest)(implicit ec: ExecutionContext): Source[SearchHit, NotUsed] =
+    Source.fromPublisher(client.publisher(searchRequest))
 
   /**
     * Create a Sink (akka stream) that create entity in ElasticSearch
     */
-  def sink[T](implicit builder: RequestBuilder[T]): Sink[T, Future[Unit]] = {
+  def sink[T](implicit builder: RequestBuilder[T], ec: ExecutionContext): Sink[T, Future[Unit]] = {
     val sinkListener = new ResponseListener[T] {
       override def onAck(resp: BulkResponseItem, original: T): Unit = ()
 
@@ -209,5 +208,5 @@ class DBConfiguration @Inject() (
     * return a new instance of DBConfiguration that points to the previous version of the index schema
     */
   def previousVersion: DBConfiguration =
-    new DBConfiguration(config, lifecycle, version - 1, ec, actorSystem)
+    new DBConfiguration(config, lifecycle, version - 1, actorSystem)
 }
diff --git a/migration/src/main/scala/org/thp/thehive/migration/th3/DBFind.scala b/migration/src/main/scala/org/thp/thehive/migration/th3/DBFind.scala
@@ -4,9 +4,9 @@ import akka.NotUsed
 import akka.stream.scaladsl.Source
 import akka.stream.stage.{AsyncCallback, GraphStage, GraphStageLogic, OutHandler}
 import akka.stream.{Attributes, Materializer, Outlet, SourceShape}
-import com.sksamuel.elastic4s.http.ElasticDsl._
-import com.sksamuel.elastic4s.http.search.{SearchHit, SearchResponse}
-import com.sksamuel.elastic4s.searches.SearchRequest
+import com.sksamuel.elastic4s.ElasticDsl._
+import com.sksamuel.elastic4s.{ElasticRequest, Show}
+import com.sksamuel.elastic4s.requests.searches.{SearchHit, SearchRequest, SearchResponse}
 import javax.inject.{Inject, Singleton}
 import org.thp.scalligraph.{InternalError, SearchError}
 import play.api.libs.json._
@@ -75,6 +75,9 @@ class DBFind(pageSize: Int, keepAlive: FiniteDuration, db: DBConfiguration, impl
     (src, total)
   }
 
+  def showQuery(request: SearchRequest): String =
+    Show[ElasticRequest].show(SearchHandler.build(request))
+
   /**
     * Search entities in ElasticSearch
     *
@@ -87,10 +90,10 @@ class DBFind(pageSize: Int, keepAlive: FiniteDuration, db: DBConfiguration, impl
   def apply(range: Option[String], sortBy: Seq[String])(query: String => SearchRequest): (Source[JsObject, NotUsed], Future[Long]) = {
     val (offset, limit) = getOffsetAndLimitFromRange(range)
     val sortDef         = DBUtils.sortDefinition(sortBy)
-    val searchRequest   = query(db.indexName).start(offset).sortBy(sortDef).version(true)
+    val searchRequest   = query(db.indexName).start(offset).sortBy(sortDef).seqNoPrimaryTerm(true)
 
     logger.debug(
-      s"search in ${searchRequest.indexesTypes.indexes.mkString(",")} / ${searchRequest.indexesTypes.types.mkString(",")} ${db.client.show(searchRequest)}"
+      s"search in ${searchRequest.indexes.values.mkString(",")} ${showQuery(searchRequest)}"
     )
     val (src, total) =
       if (limit > 2 * pageSize)
@@ -108,7 +111,7 @@ class DBFind(pageSize: Int, keepAlive: FiniteDuration, db: DBConfiguration, impl
   def apply(query: String => SearchRequest): Future[SearchResponse] = {
     val searchRequest = query(db.indexName)
     logger.debug(
-      s"search in ${searchRequest.indexesTypes.indexes.mkString(",")} / ${searchRequest.indexesTypes.types.mkString(",")} ${db.client.show(searchRequest)}"
+      s"search in ${searchRequest.indexes.values.mkString(",")} ${showQuery(searchRequest)}"
     )
 
     db.execute(searchRequest)

diff --git a/migration/src/main/scala/org/thp/thehive/migration/th3/DBGet.scala b/migration/src/main/scala/org/thp/thehive/migration/th3/DBGet.scala
@@ -1,6 +1,6 @@
 package org.thp.thehive.migration.th3
 
-import com.sksamuel.elastic4s.http.ElasticDsl._
+import com.sksamuel.elastic4s.ElasticDsl._
 import javax.inject.{Inject, Singleton}
 import org.thp.scalligraph.NotFoundError
 import play.api.libs.json.JsObject
@@ -23,7 +23,7 @@ class DBGet @Inject() (db: DBConfiguration, implicit val ec: ExecutionContext) {
       search(db.indexName)
         .query(idsQuery(id) /*.types(modelName)*/ )
         .size(1)
-        .version(true)
+        .seqNoPrimaryTerm(true)
     }.map { searchResponse =>
       searchResponse
         .hits