#1670 Link freetags to freetag taxonomy

TheHive-Project · Mar 2, 2021 · da6cd83 · da6cd83
1 parent 323b6d6
commit da6cd83
Show file tree

Hide file tree

Showing 5 changed files with 28 additions and 14 deletions.
diff --git a/ScalliGraph b/ScalliGraph
diff --git a/thehive/app/org/thp/thehive/services/OrganisationSrv.scala b/thehive/app/org/thp/thehive/services/OrganisationSrv.scala
@@ -52,7 +52,7 @@ class OrganisationSrv @Inject() (
     val activeTaxos = getByName("admin").taxonomies.toSeq
     for {
       newOrga <- createEntity(e)
-      _       <- taxonomySrv.createFreetag(newOrga)
+      _       <- taxonomySrv.createFreetagTaxonomy(newOrga)
       _       <- activeTaxos.toTry(t => organisationTaxonomySrv.create(OrganisationTaxonomy(), newOrga, t))
       _       <- auditSrv.organisation.create(newOrga, newOrga.toJson)
     } yield newOrga
@@ -187,6 +187,8 @@ object OrganisationOps {
             RichOrganisation(organisation, linkedOrganisations)
         }
 
+    def notAdmin: Traversal.V[Organisation] = traversal.hasNot(_.name, Organisation.administration.name)
+
     def isAdmin: Boolean = traversal.has(_.name, Organisation.administration.name).exists
 
     def users: Traversal.V[User] = traversal.in[RoleOrganisation].in[UserRole].v[User]

diff --git a/thehive/app/org/thp/thehive/services/TagSrv.scala b/thehive/app/org/thp/thehive/services/TagSrv.scala
@@ -6,7 +6,7 @@ import org.apache.tinkerpop.gremlin.structure.Vertex
 import org.thp.scalligraph.auth.AuthContext
 import org.thp.scalligraph.models.{Database, Entity}
 import org.thp.scalligraph.services.config.{ApplicationConfig, ConfigItem}
-import org.thp.scalligraph.services.{IntegrityCheckOps, VertexSrv}
+import org.thp.scalligraph.services.{EdgeSrv, IntegrityCheckOps, VertexSrv}
 import org.thp.scalligraph.traversal.TraversalOps._
 import org.thp.scalligraph.traversal.{Converter, Graph, Traversal}
 import org.thp.scalligraph.utils.FunctionalCondition.When
@@ -21,18 +21,17 @@ import scala.util.{Success, Try}
 @Singleton
 class TagSrv @Inject() (
     organisationSrv: OrganisationSrv,
+    taxonomySrv: TaxonomySrv,
     appConfig: ApplicationConfig,
     @Named("integrity-check-actor") integrityCheckActor: ActorRef
 ) extends VertexSrv[Tag] {
 
+  val taxonomyTagSrv = new EdgeSrv[TaxonomyTag, Taxonomy, Tag]
   private val freeTagColourConfig: ConfigItem[String, String] =
     appConfig.item[String]("tags.freeTagColour", "Default colour for free tags")
 
   def freeTagColour: String = freeTagColourConfig.get
 
-  private def freeTag(tagName: String)(implicit graph: Graph, authContext: AuthContext): Tag =
-    Tag(freeTagNamespace, tagName, None, None, freeTagColour)
-
   private def freeTagNamespace(implicit graph: Graph, authContext: AuthContext): String =
     s"_freetags_${organisationSrv.currentId(graph, authContext).value}"
 
@@ -60,9 +59,16 @@ class TagSrv @Inject() (
         startTraversal
           .getByName(freeTagNamespace, tagName, None)
           .getOrFail("Tag")
-          .orElse(create(freeTag(tagName)))
+          .orElse(createFreeTag(tagName))
       }(Success(_))
 
+  private def createFreeTag(tagName: String)(implicit graph: Graph, authContext: AuthContext): Try[Tag with Entity] =
+    for {
+      freetagTaxonomy <- taxonomySrv.getFreetagTaxonomy
+      tag             <- createEntity(Tag(freeTagNamespace, tagName, None, None, freeTagColour))
+      _               <- taxonomyTagSrv.create(TaxonomyTag(), freetagTaxonomy, tag)
+    } yield tag
+
   def create(tag: Tag)(implicit graph: Graph, authContext: AuthContext): Try[Tag with Entity] = {
     integrityCheckActor ! EntityAdded("Tag")
     super.createEntity(tag)

diff --git a/thehive/app/org/thp/thehive/services/TaxonomySrv.scala b/thehive/app/org/thp/thehive/services/TaxonomySrv.scala
@@ -31,17 +31,23 @@ class TaxonomySrv @Inject() (organisationSrv: OrganisationSrv, tagSrv: TagSrv) e
       richTaxonomy <- Try(RichTaxonomy(taxonomy, tags))
     } yield richTaxonomy
 
-  def createFreetag(organisation: Organisation with Entity)(implicit graph: Graph, authContext: AuthContext): Try[RichTaxonomy] = {
+  def createFreetagTaxonomy(organisation: Organisation with Entity)(implicit graph: Graph, authContext: AuthContext): Try[Taxonomy with Entity] = {
     val customTaxo = Taxonomy(s"_freetags_${organisation._id.value}", "Custom taxonomy", 1)
     for {
-      taxonomy     <- createEntity(customTaxo)
-      richTaxonomy <- Try(RichTaxonomy(taxonomy, Seq()))
-      _            <- organisationTaxonomySrv.create(OrganisationTaxonomy(), organisation, taxonomy)
-    } yield richTaxonomy
+      taxonomy <- createEntity(customTaxo)
+      _        <- organisationTaxonomySrv.create(OrganisationTaxonomy(), organisation, taxonomy)
+    } yield taxonomy
   }
 
+  def getFreetagTaxonomy(implicit graph: Graph, authContext: AuthContext): Try[Taxonomy with Entity] =
+    getByName(s"_freetags_${organisationSrv.currentId.value}")
+      .getOrFail("FreetagTaxonomy")
+      .orElse {
+        organisationSrv.current.notAdmin.getOrFail("Organisation").flatMap(createFreetagTaxonomy)
+      }
+
   override def getByName(name: String)(implicit graph: Graph): Traversal.V[Taxonomy] =
-    Try(startTraversal.getByNamespace(name)).getOrElse(startTraversal.limit(0))
+    startTraversal.getByNamespace(name)
 
   def update(taxonomy: Taxonomy with Entity, input: Taxonomy)(implicit graph: Graph): Try[RichTaxonomy] =
     for {

diff --git a/thehive/test/org/thp/thehive/DatabaseBuilder.scala b/thehive/test/org/thp/thehive/DatabaseBuilder.scala
@@ -155,7 +155,7 @@ class DatabaseBuilder @Inject() (
             .startTraversal
             .hasNot(_.name, "admin")
             .filterNot(_.taxonomies.freetags)
-            .foreach(o => taxonomySrv.createFreetag(o))
+            .foreach(o => taxonomySrv.createFreetagTaxonomy(o))
 
           // TODO: get tags from entity and create freetag for each
 //          // Add each tag to its Organisation's FreeTags taxonomy