#1410 Add missing properties in describe API

TheHive-Project · Jun 27, 2020 · bcb7413 · bcb7413
1 parent dc96002
commit bcb7413
Show file tree

Hide file tree

Showing 3 changed files with 56 additions and 36 deletions.
diff --git a/thehive/app/org/thp/thehive/controllers/v1/DescribeCtrl.scala b/thehive/app/org/thp/thehive/controllers/v1/DescribeCtrl.scala
@@ -12,7 +12,7 @@ import org.thp.scalligraph.services.config.ApplicationConfig.durationFormat
 import org.thp.scalligraph.services.config.{ApplicationConfig, ConfigItem}
 import org.thp.scalligraph.steps.StepsOps._
 import org.thp.scalligraph.utils.Hash
-import org.thp.thehive.services.CustomFieldSrv
+import org.thp.thehive.services.{CustomFieldSrv, ImpactStatusSrv, ResolutionStatusSrv}
 import play.api.Logger
 import play.api.cache.SyncCacheApi
 import play.api.inject.Injector
@@ -34,6 +34,8 @@ class DescribeCtrl @Inject() (
 //    logCtrl: LogCtrl,
     auditCtrl: AuditCtrl,
     customFieldSrv: CustomFieldSrv,
+    impactStatusSrv: ImpactStatusSrv,
+    resolutionStatusSrv: ResolutionStatusSrv,
     injector: Injector,
     @Named("with-thehive-schema") db: Database,
     applicationConfig: ApplicationConfig
@@ -88,12 +90,22 @@ class DescribeCtrl @Inject() (
     customFieldSrv.initSteps.toList.map(cf => PropertyDescription(s"customFields.${cf.name}", cf.`type`.toString))
   }
 
+  def impactStatus: PropertyDescription = db.roTransaction { implicit graph =>
+    PropertyDescription("impactStatus", "enumeration", impactStatusSrv.initSteps.toList.map(s => JsString(s.value)))
+  }
+
+  def resolutionStatus: PropertyDescription = db.roTransaction { implicit graph =>
+    PropertyDescription("resolutionStatus", "enumeration", resolutionStatusSrv.initSteps.toList.map(s => JsString(s.value)))
+  }
+
   def customDescription(model: String, propertyName: String): Option[Seq[PropertyDescription]] = (model, propertyName) match {
-//    case (_, "owner") => Some(Seq(PropertyDescription("owner", "user")))
-//    case ("case", "status") =>
-//      Some(
-//        Seq(PropertyDescription("status", "enumeration", Seq(JsString("Open"), JsString("Resolved"), JsString("Deleted"), JsString("Duplicated"))))
-//      )
+    case (_, "assignee") => Some(Seq(PropertyDescription("assignee", "user")))
+    case ("case", "status") =>
+      Some(
+        Seq(PropertyDescription("status", "enumeration", Seq(JsString("Open"), JsString("Resolved"), JsString("Deleted"), JsString("Duplicated"))))
+      )
+    case ("case", "impactStatus")     => Some(Seq(impactStatus))
+    case ("case", "resolutionStatus") => Some(Seq(resolutionStatus))
 //    //case ("observable", "status") =>
 //    //  Some(PropertyDescription("status", "enumeration", Seq(JsString("Ok"))))
 //    //case ("observable", "dataType") =>
@@ -116,19 +128,19 @@ class DescribeCtrl @Inject() (
 //          )
 //        )
 //      )
-//    case (_, "tlp") =>
-//      Some(Seq(PropertyDescription("tlp", "number", Seq(JsNumber(0), JsNumber(1), JsNumber(2), JsNumber(3)), Seq("white", "green", "amber", "red"))))
-//    case (_, "pap") =>
-//      Some(Seq(PropertyDescription("pap", "number", Seq(JsNumber(0), JsNumber(1), JsNumber(2), JsNumber(3)), Seq("white", "green", "amber", "red"))))
-//    case (_, "severity") =>
-//      Some(
-//        Seq(
-//          PropertyDescription("severity", "number", Seq(JsNumber(1), JsNumber(2), JsNumber(3), JsNumber(4)), Seq("low", "medium", "high", "critical"))
-//        )
-//      )
-//    case (_, "createdBy")    => Some(Seq(PropertyDescription("createdBy", "user")))
-//    case (_, "updatedBy")    => Some(Seq(PropertyDescription("updatedBy", "user")))
-//    case (_, "customFields") => Some(customFields)
+    case (_, "tlp") =>
+      Some(Seq(PropertyDescription("tlp", "number", Seq(JsNumber(0), JsNumber(1), JsNumber(2), JsNumber(3)), Seq("white", "green", "amber", "red"))))
+    case (_, "pap") =>
+      Some(Seq(PropertyDescription("pap", "number", Seq(JsNumber(0), JsNumber(1), JsNumber(2), JsNumber(3)), Seq("white", "green", "amber", "red"))))
+    case (_, "severity") =>
+      Some(
+        Seq(
+          PropertyDescription("severity", "number", Seq(JsNumber(1), JsNumber(2), JsNumber(3), JsNumber(4)), Seq("low", "medium", "high", "critical"))
+        )
+      )
+    case (_, "_createdBy")   => Some(Seq(PropertyDescription("_createdBy", "user")))
+    case (_, "_updatedBy")   => Some(Seq(PropertyDescription("_updatedBy", "user")))
+    case (_, "customFields") => Some(customFields)
 //    case ("case_artifact_job" | "action", "status") =>
 //      Some(
 //        Seq(

diff --git a/thehive/app/org/thp/thehive/controllers/v1/Properties.scala b/thehive/app/org/thp/thehive/controllers/v1/Properties.scala
@@ -1,15 +1,13 @@
 package org.thp.thehive.controllers.v1
 
-import gremlin.scala.{__, By, Key, Vertex}
 import javax.inject.{Inject, Singleton}
 import org.thp.scalligraph.BadRequestError
 import org.thp.scalligraph.controllers.FPathElem
 import org.thp.scalligraph.models.UniMapping
 import org.thp.scalligraph.query.{NoValue, PublicProperty, PublicPropertyListBuilder}
-import org.thp.scalligraph.services._
 import org.thp.scalligraph.steps.IdMapping
 import org.thp.scalligraph.steps.StepsOps._
-import org.thp.thehive.models.AlertCase
+import org.thp.thehive.models.CaseStatus
 import org.thp.thehive.services.{
   AlertSrv,
   AlertSteps,
@@ -64,19 +62,8 @@ class Properties @Inject() (
       .property("pap", UniMapping.int)(_.field.updatable)
       .property("read", UniMapping.boolean)(_.field.updatable)
       .property("follow", UniMapping.boolean)(_.field.updatable)
-      .property("status", UniMapping.string)(
-        _.select(
-          _.project(
-            _.apply(By(Key[Boolean]("read")))
-              .and(By(__[Vertex].outToE[AlertCase].limit(1).count()))
-          ).map {
-            case (true, caseCount) if caseCount == 0L  => "Ignored"
-            case (true, caseCount) if caseCount == 1L  => "New"
-            case (false, caseCount) if caseCount == 0L => "Ignored"
-            case (false, caseCount) if caseCount == 1L => "Imported"
-          }
-        ).readonly
-      )
+      .property("read", UniMapping.boolean)(_.field.updatable)
+      .property("imported", UniMapping.boolean)(_.select(_.imported).readonly)
       .property("summary", UniMapping.string.optional)(_.field.updatable)
       .property("user", UniMapping.string)(_.field.updatable)
       .property("customFields", UniMapping.identity[JsValue])(_.subSelect {
@@ -124,7 +111,7 @@ class Properties @Inject() (
       .property("flag", UniMapping.boolean)(_.field.updatable)
       .property("tlp", UniMapping.int)(_.field.updatable)
       .property("pap", UniMapping.int)(_.field.updatable)
-      .property("status", UniMapping.string)(_.field.updatable)
+      .property("status", UniMapping.enum(CaseStatus))(_.field.updatable)
       .property("summary", UniMapping.string.optional)(_.field.updatable)
       .property("assignee", UniMapping.string.optional)(_.select(_.user.login).custom { (_, login, vertex, _, graph, authContext) =>
         for {
@@ -136,6 +123,25 @@ class Properties @Inject() (
           }
         } yield Json.obj("owner" -> user.map(_.login))
       })
+      .property("impactStatus", UniMapping.string.optional)(_.select(_.impactStatus.value).custom { (_, value, vertex, _, graph, authContext) =>
+        caseSrv
+          .get(vertex)(graph)
+          .getOrFail("Case")
+          .flatMap { c =>
+            value.fold(caseSrv.unsetImpactStatus(c)(graph, authContext))(caseSrv.setImpactStatus(c, _)(graph, authContext))
+          }
+          .map(_ => Json.obj("impactStatus" -> value))
+      })
+      .property("resolutionStatus", UniMapping.string.optional)(_.select(_.resolutionStatus.value).custom {
+        (_, value, vertex, _, graph, authContext) =>
+          caseSrv
+            .get(vertex)(graph)
+            .getOrFail("Case")
+            .flatMap { c =>
+              value.fold(caseSrv.unsetResolutionStatus(c)(graph, authContext))(caseSrv.setResolutionStatus(c, _)(graph, authContext))
+            }
+            .map(_ => Json.obj("resolutionStatus" -> value))
+      })
       .build
 
   lazy val caseTemplate: List[PublicProperty[_, _]] =

diff --git a/thehive/app/org/thp/thehive/services/AlertSrv.scala b/thehive/app/org/thp/thehive/services/AlertSrv.scala
@@ -325,6 +325,8 @@ class AlertSteps(raw: GremlinScala[Vertex])(implicit @Named("with-thehive-schema
         .has("login", authContext.userId)
     )
 
+  def imported: Traversal[Boolean, Boolean] = this.outToE[AlertCase].count.map(_ > 0)
+
   def alertUserOrganisation(
       permission: Permission
   )(implicit authContext: AuthContext): Traversal[(RichAlert, Organisation with Entity), (RichAlert, Organisation with Entity)] = {