Fixed tag colour issue

TheHive-Project · Jan 29, 2021 · b7109f1 · b7109f1
1 parent 54ae927
commit b7109f1
Show file tree

Hide file tree

Showing 6 changed files with 113 additions and 114 deletions.
diff --git a/dto/src/main/scala/org/thp/thehive/dto/v1/Tag.scala b/dto/src/main/scala/org/thp/thehive/dto/v1/Tag.scala
@@ -3,11 +3,11 @@ package org.thp.thehive.dto.v1
 import play.api.libs.json.{Json, OFormat}
 
 case class OutputTag(
-  namespace: String,
-  predicate: String,
-  value: Option[String],
-  description: Option[String],
-  colour: String
+    namespace: String,
+    predicate: String,
+    value: Option[String],
+    description: Option[String],
+    colour: String
 )
 
 object OutputTag {

diff --git a/dto/src/main/scala/org/thp/thehive/dto/v1/Taxonomy.scala b/dto/src/main/scala/org/thp/thehive/dto/v1/Taxonomy.scala
@@ -1,45 +1,47 @@
 package org.thp.thehive.dto.v1
 
+import play.api.libs.json.Json.WithDefaultValues
+
 import java.util.Date
 import play.api.libs.json.{JsObject, Json, OFormat}
 
 /*
 Format based on :
 https://tools.ietf.org/id/draft-dulaunoy-misp-taxonomy-format-04.html
-*/
+ */
 
 case class InputTaxonomy(
-  namespace: String,
-  description: String,
-  version: Int,
-  `type`: Option[Seq[String]],
-  exclusive: Option[Boolean],
-  predicates: Seq[InputPredicate],
-  values: Option[Seq[InputValue]]
+    namespace: String,
+    description: String,
+    version: Int,
+    exclusive: Option[Boolean],
+    predicates: Seq[InputPredicate],
+    values: Seq[InputValue] = Nil
 )
 
 case class InputPredicate(
-  value: String,
-  expanded: Option[String],
-  exclusive: Option[Boolean],
-  description: Option[String]
+    value: String,
+    expanded: Option[String],
+    exclusive: Option[Boolean],
+    description: Option[String],
+    colour: Option[String]
 )
 
 case class InputValue(
-  predicate: String,
-  entry: Seq[InputEntry]
+    predicate: String,
+    entry: Seq[InputEntry]
 )
 
 case class InputEntry(
-  value: String,
-  expanded: Option[String],
-  colour: Option[String],
-  description: Option[String],
-  numerical_value: Option[Int]
+    value: String,
+    expanded: Option[String],
+    colour: Option[String],
+    description: Option[String],
+    numerical_value: Option[Int]
 )
 
 object InputTaxonomy {
-  implicit val format: OFormat[InputTaxonomy] = Json.format[InputTaxonomy]
+  implicit val format: OFormat[InputTaxonomy] = Json.configured[WithDefaultValues].format[InputTaxonomy]
 }
 
 object InputPredicate {
@@ -55,17 +57,17 @@ object InputEntry {
 }
 
 case class OutputTaxonomy(
-  _id: String,
-  _type: String,
-  _createdBy: String,
-  _updatedBy: Option[String] = None,
-  _createdAt: Date,
-  _updatedAt: Option[Date] = None,
-  namespace: String,
-  description: String,
-  version: Int,
-  tags: Seq[OutputTag],
-  extraData: JsObject
+    _id: String,
+    _type: String,
+    _createdBy: String,
+    _updatedBy: Option[String] = None,
+    _createdAt: Date,
+    _updatedAt: Option[Date] = None,
+    namespace: String,
+    description: String,
+    version: Int,
+    tags: Seq[OutputTag],
+    extraData: JsObject
 )
 
 object OutputTaxonomy {

diff --git a/thehive/app/org/thp/thehive/controllers/v1/Properties.scala b/thehive/app/org/thp/thehive/controllers/v1/Properties.scala
@@ -472,7 +472,8 @@ class Properties @Inject() (
     val namespace = UMapping.string.getProperty(v, "namespace")
     val predicate = UMapping.string.getProperty(v, "predicate")
     val value     = UMapping.string.optional.getProperty(v, "value")
-    Tag(namespace, predicate, value, None, "#000000").toString
+    val colour    = UMapping.string.optional.getProperty(v, "colour")
+    Tag(namespace, predicate, value, None, colour.getOrElse("#000000")).toString
   }
 
 }
diff --git a/thehive/app/org/thp/thehive/controllers/v1/TaxonomyCtrl.scala b/thehive/app/org/thp/thehive/controllers/v1/TaxonomyCtrl.scala
@@ -1,6 +1,5 @@
 package org.thp.thehive.controllers.v1
 
-import javax.inject.{Inject, Named}
 import net.lingala.zip4j.ZipFile
 import net.lingala.zip4j.model.FileHeader
 import org.apache.tinkerpop.gremlin.structure.Graph
@@ -19,23 +18,23 @@ import org.thp.thehive.services.{TagSrv, TaxonomySrv}
 import play.api.libs.json.{JsArray, Json}
 import play.api.mvc.{Action, AnyContent, Results}
 
+import javax.inject.{Inject, Named}
 import scala.collection.JavaConverters._
 import scala.util.{Failure, Success, Try}
 
 class TaxonomyCtrl @Inject() (
-  entrypoint: Entrypoint,
-  properties: Properties,
-  taxonomySrv: TaxonomySrv,
-  tagSrv: TagSrv,
-  @Named("with-thehive-schema") implicit val db: Database
-) extends QueryableCtrl with TaxonomyRenderer {
+    entrypoint: Entrypoint,
+    properties: Properties,
+    taxonomySrv: TaxonomySrv,
+    tagSrv: TagSrv,
+    @Named("with-thehive-schema") implicit val db: Database
+) extends QueryableCtrl
+    with TaxonomyRenderer {
 
-  override val entityName: String = "taxonomy"
+  override val entityName: String                 = "taxonomy"
   override val publicProperties: PublicProperties = properties.taxonomy
   override val initialQuery: Query =
-    Query.init[Traversal.V[Taxonomy]]("listTaxonomy", (graph, authContext) =>
-      taxonomySrv.startTraversal(graph).visible(authContext)
-    )
+    Query.init[Traversal.V[Taxonomy]]("listTaxonomy", (graph, authContext) => taxonomySrv.startTraversal(graph).visible(authContext))
   override val getQuery: ParamQuery[EntityIdOrName] =
     Query.initWithParam[EntityIdOrName, Traversal.V[Taxonomy]](
       "getTaxonomy",
@@ -44,19 +43,17 @@ class TaxonomyCtrl @Inject() (
     )
   override val pageQuery: ParamQuery[OutputParam] =
     Query.withParam[OutputParam, Traversal.V[Taxonomy], IteratorOutput](
-    "page",
-    FieldsParser[OutputParam],
-    {
-      case (OutputParam(from, to, extraData), taxoSteps, authContext) =>
-        taxoSteps.richPage(from, to, extraData.contains("total")) {
-          _.richTaxonomyWithCustomRenderer(taxoStatsRenderer(extraData - "total"))
-        }
-    }
-  )
+      "page",
+      FieldsParser[OutputParam],
+      {
+        case (OutputParam(from, to, extraData), taxoSteps, authContext) =>
+          taxoSteps.richPage(from, to, extraData.contains("total")) {
+            _.richTaxonomyWithCustomRenderer(taxoStatsRenderer(extraData - "total"))
+          }
+      }
+    )
   override val outputQuery: Query =
-    Query.outputWithContext[RichTaxonomy, Traversal.V[Taxonomy]]((traversal, _) =>
-    traversal.richTaxonomy
-  )
+    Query.outputWithContext[RichTaxonomy, Traversal.V[Taxonomy]]((traversal, _) => traversal.richTaxonomy)
   override val extraQueries: Seq[ParamQuery[_]] = Seq(
     Query[Traversal.V[Taxonomy], Traversal.V[Tag]]("tags", (traversal, _) => traversal.tags)
   )
@@ -68,24 +65,25 @@ class TaxonomyCtrl @Inject() (
         for {
           richTaxonomy <- createFromInput(request.body("taxonomy"))
         } yield Results.Created(richTaxonomy.toJson)
-  }
+      }
 
   def importZip: Action[AnyContent] =
     entrypoint("import taxonomies zip")
       .extract("file", FieldsParser.file.on("file"))
       .authPermitted(Permissions.manageTaxonomy) { implicit request =>
         val file: FFile = request.body("file")
-        val zipFile = new ZipFile(file.filepath.toString)
+        val zipFile     = new ZipFile(file.filepath.toString)
         val headers = zipFile
           .getFileHeaders
           .iterator()
           .asScala
 
         for {
-          inputTaxos <- headers
-            .filter(h => h.getFileName.endsWith("machinetag.json"))
-            .toTry(parseJsonFile(zipFile, _))
-          richTaxos = inputTaxos.foldLeft[JsArray](JsArray.empty)((array, taxo) => {
+          inputTaxos <-
+            headers
+              .filter(h => h.getFileName.endsWith("machinetag.json"))
+              .toTry(parseJsonFile(zipFile, _))
+          richTaxos = inputTaxos.foldLeft[JsArray](JsArray.empty) { (array, taxo) =>
             val res = db.tryTransaction { implicit graph =>
               createFromInput(taxo)
             } match {
@@ -95,35 +93,35 @@ class TaxonomyCtrl @Inject() (
                 Json.obj("namespace" -> t.namespace, "status" -> "Success", "tagsImported" -> t.tags.size)
             }
             array :+ res
-          })
+          }
         } yield Results.Created(richTaxos)
       }
 
-  private def parseJsonFile(zipFile: ZipFile, h: FileHeader): Try[InputTaxonomy] = {
+  private def parseJsonFile(zipFile: ZipFile, h: FileHeader): Try[InputTaxonomy] =
     Try(Json.parse(zipFile.getInputStream(h)).as[InputTaxonomy]).recoverWith {
       case _ => Failure(BadRequestError(s"File '${h.getFileName}' does not comply with the MISP taxonomy formatting"))
     }
-  }
 
   private def createFromInput(inputTaxo: InputTaxonomy)(implicit graph: Graph, authContext: AuthContext): Try[RichTaxonomy] = {
     // Create tags
-    val tagValues = inputTaxo.values.getOrElse(Seq())
-    val tags = tagValues.flatMap(value => {
-      value.entry.map(e =>
-        Tag(inputTaxo.namespace,
-          value.predicate,
-          Some(e.value),
-          e.expanded,
-          e.colour.getOrElse(tagSrv.defaultColour)
-        )
-      )
-    })
+    val predicatesWithValue = inputTaxo.values.map(_.predicate).distinct
+    val predicateWithNoTags = inputTaxo.predicates.filterNot(p => predicatesWithValue.contains(p.value))
 
+    val tags = inputTaxo.values.flatMap { value =>
+      value
+        .entry
+        .map { e =>
+          Tag(
+            inputTaxo.namespace,
+            value.predicate,
+            Some(e.value),
+            e.expanded,
+            e.colour.getOrElse(tagSrv.defaultColour)
+          )
+        }
+    }
     // Create a tag for predicates with no tags associated
-    val predicateWithNoTags = inputTaxo.predicates.map(_.value).diff(tagValues.map(_.predicate))
-    val allTags = tags ++ predicateWithNoTags.map(p =>
-      Tag(inputTaxo.namespace, p, None, None, tagSrv.defaultColour)
-    )
+    val allTags = tags ++ predicateWithNoTags.map(p => Tag(inputTaxo.namespace, p.value, None, None, p.colour.getOrElse(tagSrv.defaultColour)))
 
     if (inputTaxo.namespace.isEmpty)
       Failure(BadRequestError(s"A taxonomy with no namespace cannot be imported"))
@@ -160,13 +158,14 @@ class TaxonomyCtrl @Inject() (
     entrypoint("delete taxonomy")
       .authPermittedTransaction(db, Permissions.manageTaxonomy) { implicit request => implicit graph =>
         for {
-        taxo <- taxonomySrv
-          .get(EntityIdOrName(taxoId))
-          .visible
-          .getOrFail("Taxonomy")
-        tags <- Try(taxonomySrv.get(taxo).tags.toSeq)
-        _ <- tags.toTry(t => tagSrv.delete(t))
-        _ <- taxonomySrv.delete(taxo)
+          taxo <-
+            taxonomySrv
+              .get(EntityIdOrName(taxoId))
+              .visible
+              .getOrFail("Taxonomy")
+          tags <- Try(taxonomySrv.get(taxo).tags.toSeq)
+          _    <- tags.toTry(t => tagSrv.delete(t))
+          _    <- taxonomySrv.delete(taxo)
         } yield Results.NoContent
       }
 

diff --git a/thehive/test/org/thp/thehive/controllers/v1/TaxonomyCtrlTest.scala b/thehive/test/org/thp/thehive/controllers/v1/TaxonomyCtrlTest.scala
@@ -10,10 +10,10 @@ import play.api.mvc.{AnyContentAsMultipartFormData, MultipartFormData}
 import play.api.test.{FakeRequest, PlaySpecification}
 
 case class TestTaxonomy(
-  namespace: String,
-  description: String,
-  version: Int,
-  tags: List[OutputTag]
+    namespace: String,
+    description: String,
+    version: Int,
+    tags: List[OutputTag]
 )
 
 object TestTaxonomy {
@@ -34,20 +34,17 @@ class TaxonomyCtrlTest extends PlaySpecification with TestAppBuilder {
       "A test taxonomy",
       1,
       None,
-      None,
       List(
-        InputPredicate("pred1", None, None, None),
-        InputPredicate("pred2", None, None, None)
+        InputPredicate("pred1", None, None, None, None),
+        InputPredicate("pred2", None, None, None, None)
       ),
-      Some(
-        List(
-          InputValue("pred1", List(InputEntry("entry1", None, None, None, None))),
-          InputValue(
-            "pred2",
-            List(
-              InputEntry("entry2", None, None, None, None),
-              InputEntry("entry21", None, None, None, None)
-            )
+      List(
+        InputValue("pred1", List(InputEntry("entry1", None, Some("#ffa800"), None, None))),
+        InputValue(
+          "pred2",
+          List(
+            InputEntry("entry2", None, Some("#00ad1c"), None, None),
+            InputEntry("entry21", None, Some("#00ad1c"), None, None)
           )
         )
       )
@@ -61,16 +58,16 @@ class TaxonomyCtrlTest extends PlaySpecification with TestAppBuilder {
       val result = app[TaxonomyCtrl].create(request)
       status(result) must beEqualTo(201).updateMessage(s => s"$s\n${contentAsString(result)}")
 
-      val resultCase = contentAsJson(result).as[OutputTaxonomy]
+      val resultTaxo = contentAsJson(result).as[OutputTaxonomy]
 
-      TestTaxonomy(resultCase) must_=== TestTaxonomy(
+      TestTaxonomy(resultTaxo) must_=== TestTaxonomy(
         "test-taxo",
         "A test taxonomy",
         1,
         List(
-          OutputTag("test-taxo", "pred1", Some("entry1"), None, "#000000"),
-          OutputTag("test-taxo", "pred2", Some("entry2"), None, "#000000"),
-          OutputTag("test-taxo", "pred2", Some("entry21"), None, "#000000")
+          OutputTag("test-taxo", "pred1", Some("entry1"), None, "#ffa800"),
+          OutputTag("test-taxo", "pred2", Some("entry2"), None, "#00ad1c"),
+          OutputTag("test-taxo", "pred2", Some("entry21"), None, "#00ad1c")
         )
       )
     }
@@ -123,7 +120,7 @@ class TaxonomyCtrlTest extends PlaySpecification with TestAppBuilder {
         "taxonomy1",
         "The taxonomy 1",
         1,
-        List(OutputTag("taxonomy1", "pred1", Some("value1"), None, "#000000"))
+        List(OutputTag("taxonomy1", "pred1", Some("value1"), None, "#00f300"))
       )
     }
 

diff --git a/thehive/test/resources/data/Tag.json b/thehive/test/resources/data/Tag.json
@@ -74,6 +74,6 @@
     "namespace": "taxonomy1",
     "predicate": "pred1",
     "value": "value1",
-    "colour": "#000000"
+    "colour": "#00f300"
   }
 ]