#237 Migrate alerts to store attachment in datastore

Author: To-om

thehive-backend/app/models/Migration.scala

@@ -3,14 +3,18 @@ package models
 import java.util.Date
 import javax.inject.Inject
 
+import akka.NotUsed
 import akka.stream.Materializer
+import akka.stream.scaladsl.Source
 import org.elastic4play.models.BaseModelDef
 import org.elastic4play.services._
+import org.elastic4play.services.JsonFormat.attachmentFormat
 import org.elastic4play.utils
 import org.elastic4play.utils.{ Hasher, RichJson }
-import play.api.{ Configuration, Logger }
 import play.api.libs.json.JsValue.jsValueToJsLookup
 import play.api.libs.json._
+import play.api.{ Configuration, Logger }
+import services.AlertSrv
 
 import scala.collection.immutable.{ Set ⇒ ISet }
 import scala.concurrent.{ ExecutionContext, Future }
@@ -21,6 +25,9 @@ case class UpdateMispAlertArtifact() extends EventMessage
 
 class Migration(
     mispCaseTemplate: Option[String],
+    mainHash: String,
+    extraHashes: Seq[String],
+    datastoreName: String,
     models: ISet[BaseModelDef],
     dblists: DBLists,
     eventSrv: EventSrv,
@@ -33,10 +40,17 @@ class Migration(
     eventSrv: EventSrv,
     ec: ExecutionContext,
     materializer: Materializer) = {
-    this(configuration.getString("misp.caseTemplate"), models, dblists, eventSrv, ec, materializer)
+    this(
+      configuration.getString("misp.caseTemplate"),
+      configuration.getString("datastore.hash.main").get,
+      configuration.getStringSeq("datastore.hash.extra").get,
+      configuration.getString("datastore.name").get,
+      models, dblists,
+      eventSrv, ec, materializer)
   }
 
   import org.elastic4play.services.Operation._
+
   val logger = Logger(getClass)
   private var requireUpdateMispAlertArtifact = false
 
@@ -126,10 +140,65 @@ class Migration(
               "follow" → (misp \ "follow").as[JsBoolean])
         },
         removeEntity("audit")(o ⇒ (o \ "objectType").asOpt[String].contains("alert")))
-    case DatabaseState(9) ⇒ Nil
+    case ds @ DatabaseState(9) ⇒
+      object Base64 {
+        def unapply(data: String): Option[Array[Byte]] = Try(java.util.Base64.getDecoder.decode(data)).toOption
+      }
+
+      var dataIds = Set.empty[String]
+      def containsOrAdd(id: String) = {
+        dataIds.synchronized {
+          if (dataIds.contains(id)) true
+          else {
+            dataIds = dataIds + id
+            false
+          }
+        }
+      }
+      val mainHasher = Hasher(mainHash)
+      val extraHashers = Hasher(mainHash +: extraHashes: _*)
+      Seq(
+        Operation((f: String ⇒ Source[JsObject, NotUsed]) ⇒ {
+        case "alert" ⇒ f("alert").flatMapConcat { alert ⇒
+          val artifactsAndData = Future.traverse((alert \ "artifacts").asOpt[List[JsObject]].getOrElse(Nil)) { artifact ⇒
+            (artifact \ "data").asOpt[String]
+              .collect {
+                case AlertSrv.dataExtractor(filename, contentType, data @ Base64(rawData)) ⇒
+                  val attachmentId = mainHasher.fromByteArray(rawData).head.toString()
+                  ds.getEntity(datastoreName, s"${attachmentId}_0")
+                    .map(_ ⇒ Nil)
+                    .recover {
+                      case _ if containsOrAdd(attachmentId) ⇒ Nil
+                      case _ ⇒
+                        Seq(Json.obj(
+                          "_type" → datastoreName,
+                          "_id" → s"${attachmentId}_0",
+                          "data" → data))
+                    }
+                    .map { dataEntity ⇒
+                      val attachment = Attachment(filename, extraHashers.fromByteArray(rawData), rawData.length.toLong, contentType, attachmentId)
+                      (artifact - "data" + ("attachment" → Json.toJson(attachment))) → dataEntity
+                    }
+              }
+              .getOrElse(Future.successful(artifact → Nil))
+          }
+          Source.fromFuture(artifactsAndData)
+            .mapConcat { ad ⇒
+              val updatedAlert = alert + ("artifacts" → JsArray(ad.map(_._1)))
+              updatedAlert :: ad.flatMap(_._2)
+            }
+        }
+        case other ⇒ f(other)
+      }),
+        mapAttribute("alert", "status") {
+          case JsString("Update") ⇒ JsString("Updated")
+          case JsString("Ignore") ⇒ JsString("Ignored")
+          case other              ⇒ other
+        })
   }
 
   private val requestCounter = new java.util.concurrent.atomic.AtomicInteger(0)
+
   def getRequestId: String = {
     utils.Instance.id + ":mig:" + requestCounter.incrementAndGet()
   }

thehive-backend/app/services/AlertSrv.scala

@@ -28,6 +28,10 @@ trait AlertTransformer {
 
 case class CaseSimilarity(caze: Case, similarIOCCount: Int, iocCount: Int, similarArtifactCount: Int, artifactCount: Int)
 
+object AlertSrv {
+  val dataExtractor = "^(.*);(.*);(.*)".r
+}
+
 class AlertSrv(
     templates: Map[String, String],
     alertModel: AlertModel,
@@ -77,6 +81,7 @@ class AlertSrv(
     mat)
 
   private[AlertSrv] lazy val logger = Logger(getClass)
+  import AlertSrv._
 
   def create(fields: Fields)(implicit authContext: AuthContext): Future[Alert] = {
 
@@ -149,8 +154,6 @@ class AlertSrv(
       .recover { case _ ⇒ None }
   }
 
-  private val dataExtractor = "^(.*);(.*);(.*)".r
-
   def createCase(alert: Alert, customCaseTemplate: Option[String])(implicit authContext: AuthContext): Future[Case] = {
     alert.caze() match {
       case Some(id) ⇒ caseSrv.get(id)