#2206 Rewrite using scala style

TheHive-Project · Oct 1, 2021 · 56caa31 · 56caa31
1 parent 2b6550a
commit 56caa31
Show file tree

Hide file tree

Showing 6 changed files with 27 additions and 41 deletions.
diff --git a/cortex/connector/src/main/scala/org/thp/thehive/connector/cortex/services/JobSrv.scala b/cortex/connector/src/main/scala/org/thp/thehive/connector/cortex/services/JobSrv.scala
@@ -195,7 +195,7 @@ class JobSrv @Inject() (
         val tags = cortexJob.report.fold[Seq[ReportTag]](Nil)(_.summary.map(_.toAnalyzerTag(job.workerName)))
         for {
           observable <- get(job).observable.getOrFail("Observable")
-          _          <- reportTagSrv.updateTags(observable, job.workerName, tags)
+          _          <- reportTagSrv.updateTags(observable, tags)
         } yield ()
       }
     else Success(())

diff --git a/migration/src/main/scala/org/thp/thehive/migration/Migrate.scala b/migration/src/main/scala/org/thp/thehive/migration/Migrate.scala
@@ -60,7 +60,7 @@ object Migrate extends App with MigrationOps {
         .valueName("http://ip1:port,ip2:port")
         .text("TheHive3 ElasticSearch URI")
         .action((u, c) => addConfig(c, "input.search.uri", u)),
-      opt[String]('i', "es-index")
+      opt[String]('e', "es-index")
         .valueName("<index>")
         .text("TheHive3 ElasticSearch index name")
         .action((i, c) => addConfig(c, "input.search.index", i)),

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
@@ -7,5 +7,5 @@ case class InputObservable(
     observable: Observable,
     organisations: Set[String],
     dataOrAttachment: Either[String, InputAttachment],
-    reportTags: List[ReportTag]
+    reportTags: Seq[ReportTag]
 )
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
@@ -12,8 +12,9 @@ import org.thp.thehive.models._
 import play.api.libs.functional.syntax._
 import play.api.libs.json.JsValue.jsValueToJsLookup
 import play.api.libs.json._
+
 import java.util.{Base64, Date}
-import scala.collection.mutable
+import scala.util.Try
 
 case class Attachment(name: String, hashes: Seq[Hash], size: Long, contentType: String, id: String)
 trait Conversion {
@@ -102,38 +103,22 @@ trait Conversion {
     )
   }
 
-  def getTaxonomies(input: Map[String, JsValue]): List[ReportTag] = {
-    val taxonomies = mutable.MutableList[ReportTag]()
-    input.foreach{ case (origin: String, value: JsValue) =>
-      (jsValueToJsLookup(value) \ "taxonomies").asOpt[List[Map[String, JsValue]]].getOrElse(List.empty).foreach(taxonomy => {
-        taxonomies += ReportTag(
-          origin,
-          taxonomy.get("level") match {
-            case Some(x) => x.asOpt[String].getOrElse("") match {
-              case "malicious" => ReportTagLevel.malicious
-              case "suspicious" => ReportTagLevel.suspicious
-              case "safe" => ReportTagLevel.safe
-              case "info" => ReportTagLevel.info
-            }
-            case None    => throw new Exception
-          },
-          taxonomy.get("namespace") match {
-            case Some(x) => x.asOpt[String].getOrElse("")
-            case None    => throw new Exception
-          },
-          taxonomy.get("predicate") match {
-            case Some(x) => x.asOpt[String].getOrElse("")
-            case None    => throw new Exception
-          },
-          taxonomy.get("value") match {
-            case Some(x) => x
-            case None    => throw new Exception
+  implicit val taxonomiesReads: Reads[Seq[ReportTag]] = Reads.JsObjectReads.map { input =>
+    input.fields.flatMap {
+      case (origin, value) =>
+        (value \ "taxonomies")
+          .asOpt[Seq[JsValue]]
+          .getOrElse(Nil)
+          .flatMap { taxonomy =>
+            for {
+              namespace <- (taxonomy \ "namespace").asOpt[String]
+              predicate <- (taxonomy \ "predicate").asOpt[String]
+              value = (taxonomy \ "value").getOrElse(JsNull)
+              levelName <- (taxonomy \ "level").asOpt[String]
+              level     <- Try(ReportTagLevel.withName(levelName)).toOption
+            } yield ReportTag(origin, level, namespace, predicate, value)
           }
-        )
-      })
     }
-
-    return taxonomies.toList
   }
 
   implicit val observableReads: Reads[InputObservable] = Reads[InputObservable] { json =>
@@ -144,9 +129,8 @@ trait Conversion {
       ioc      <- (json \ "ioc").validate[Boolean]
       sighted  <- (json \ "sighted").validate[Boolean]
       dataType <- (json \ "dataType").validate[String]
-
       tags = (json \ "tags").asOpt[Set[String]].getOrElse(Set.empty)
-      taxonomiesList = getTaxonomies(Json.parse((json \ "reports").asOpt[String].getOrElse("{}")).as[Map[String, JsValue]])
+      taxonomiesList <- Json.parse((json \ "reports").asOpt[String].getOrElse("{}")).validate[Seq[ReportTag]]
       dataOrAttachment <-
         (json \ "data")
           .validate[String]
@@ -297,7 +281,7 @@ trait Conversion {
         ),
         Set(mainOrganisation),
         dataOrAttachment,
-        List()
+        Nil
       )
 
     }
@@ -524,7 +508,7 @@ trait Conversion {
         ),
         Set(mainOrganisation),
         dataOrAttachment,
-        List()
+        Nil
       )
     }
 

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
@@ -698,7 +698,7 @@ class Output @Inject() (
       for {
         organisations  <- inputObservable.organisations.toTry(getOrganisation)
         richObservable <- createObservable(caseId, inputObservable, organisations.map(_._id).toSet)
-        _              <- reportTagSrv.updateTags(richObservable, "not_existing_origin", inputObservable.reportTags)
+        _              <- reportTagSrv.updateTags(richObservable, inputObservable.reportTags)
         case0          <- getCase(caseId)
         _              <- organisations.toTry(o => shareSrv.shareObservable(RichObservable(richObservable, None, None, Nil), case0, o._id))
       } yield IdMapping(inputObservable.metaData.id, richObservable._id)

diff --git a/thehive/app/org/thp/thehive/services/ReportTagSrv.scala b/thehive/app/org/thp/thehive/services/ReportTagSrv.scala
@@ -17,11 +17,13 @@ import scala.util.Try
 class ReportTagSrv @Inject() (observableSrv: ObservableSrv) extends VertexSrv[ReportTag] {
   val observableReportTagSrv = new EdgeSrv[ObservableReportTag, Observable, ReportTag]
 
-  def updateTags(observable: Observable with Entity, origin: String, reportTags: Seq[ReportTag])(implicit
+  def updateTags(observable: Observable with Entity, reportTags: Seq[ReportTag])(implicit
       graph: Graph,
       authContext: AuthContext
   ): Try[Unit] = {
-    observableSrv.get(observable).reportTags.fromOrigin(origin).remove()
+    reportTags.map(_.origin).distinct.foreach { origin =>
+      observableSrv.get(observable).reportTags.fromOrigin(origin).remove()
+    }
     reportTags
       .toTry { tag =>
         createEntity(tag).flatMap(t => observableReportTagSrv.create(ObservableReportTag(), observable, t))