From f087c2fa23c07cc81877ecb480c96126287de193 Mon Sep 17 00:00:00 2001
From: To-om <toom@thehive-project.org>
Date: Fri, 11 Dec 2020 18:19:52 +0100
Subject: [PATCH] #1708 Add serializer for flow messages

---
 .../thehive/controllers/v0/AuditCtrl.scala    |  1 -
 .../org/thp/thehive/services/FlowActor.scala  | 10 +++---
 .../thp/thehive/services/FlowSerializer.scala | 32 +++++++++++++++++++
 3 files changed, 36 insertions(+), 7 deletions(-)
 create mode 100644 thehive/app/org/thp/thehive/services/FlowSerializer.scala

diff --git a/thehive/app/org/thp/thehive/controllers/v0/AuditCtrl.scala b/thehive/app/org/thp/thehive/controllers/v0/AuditCtrl.scala
index 1e14c0873f..9ced6e86b6 100644
--- a/thehive/app/org/thp/thehive/controllers/v0/AuditCtrl.scala
+++ b/thehive/app/org/thp/thehive/controllers/v0/AuditCtrl.scala
@@ -13,7 +13,6 @@ import org.thp.scalligraph.traversal.{IteratorOutput, Traversal}
 import org.thp.thehive.controllers.v0.Conversion._
 import org.thp.thehive.models.{Audit, RichAudit}
 import org.thp.thehive.services.AuditOps._
-import org.thp.thehive.services.FlowActor.{AuditIds, FlowId}
 import org.thp.thehive.services._
 import play.api.libs.json.{JsArray, JsObject, Json}
 import play.api.mvc.{Action, AnyContent, Results}
diff --git a/thehive/app/org/thp/thehive/services/FlowActor.scala b/thehive/app/org/thp/thehive/services/FlowActor.scala
index 5c3411498a..8d023071cd 100644
--- a/thehive/app/org/thp/thehive/services/FlowActor.scala
+++ b/thehive/app/org/thp/thehive/services/FlowActor.scala
@@ -21,15 +21,13 @@ import play.api.cache.SyncCacheApi
 
 import scala.concurrent.duration.FiniteDuration
 
-object FlowActor {
-  case class FlowId(organisation: EntityIdOrName, caseId: Option[EntityIdOrName]) {
-    override def toString: String = s"$organisation;${caseId.getOrElse("-")}"
-  }
-  case class AuditIds(ids: Seq[EntityId])
+sealed trait FlowMessage
+case class FlowId(organisation: EntityIdOrName, caseId: Option[EntityIdOrName]) extends FlowMessage {
+  override def toString: String = s"$organisation;${caseId.getOrElse("-")}"
 }
+case class AuditIds(ids: Seq[EntityId]) extends FlowMessage
 
 class FlowActor extends Actor {
-  import FlowActor._
 
   lazy val injector: Injector           = GuiceAkkaExtension(context.system).injector
   lazy val cache: SyncCacheApi          = injector.getInstance(classOf[SyncCacheApi])
diff --git a/thehive/app/org/thp/thehive/services/FlowSerializer.scala b/thehive/app/org/thp/thehive/services/FlowSerializer.scala
new file mode 100644
index 0000000000..9cff3137d6
--- /dev/null
+++ b/thehive/app/org/thp/thehive/services/FlowSerializer.scala
@@ -0,0 +1,32 @@
+package org.thp.thehive.services
+
+import akka.serialization.Serializer
+import org.thp.scalligraph.{EntityId, EntityIdOrName}
+
+import java.io.NotSerializableException
+
+class FlowSerializer extends Serializer {
+  override def identifier: Int = -1165729876
+
+  override def includeManifest: Boolean = false
+
+  override def toBinary(o: AnyRef): Array[Byte] =
+    o match {
+      case FlowId(organisation, None)         => 0.toByte +: organisation.toString.getBytes
+      case FlowId(organisation, Some(caseId)) => 1.toByte +: s"$organisation|$caseId".getBytes
+      case AuditIds(ids)                      => 2.toByte +: ids.map(_.value).mkString("|").getBytes
+      case _                                  => throw new NotSerializableException
+    }
+
+  override def fromBinary(bytes: Array[Byte], manifest: Option[Class[_]]): AnyRef =
+    bytes(0) match {
+      case 0 => FlowId(EntityIdOrName(new String(bytes.tail)), None)
+      case 1 =>
+        new String(bytes.tail).split('|') match {
+          case Array(organisation, caseId) => FlowId(EntityIdOrName(organisation), Some(EntityIdOrName(caseId)))
+          case _                           => throw new NotSerializableException
+        }
+      case 2 => AuditIds(new String(bytes.tail).split('|').toSeq.map(EntityId.apply))
+      case _ => throw new NotSerializableException
+    }
+}