#10 Add parameter to AuxSrv to filter out unaudited attributes

TheHive-Project · Feb 21, 2017 · 3ad86fc · 3ad86fc
1 parent 40df9ec
commit 3ad86fc
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 9 deletions.
diff --git a/app/org/elastic4play/controllers/Renderer.scala b/app/org/elastic4play/controllers/Renderer.scala
@@ -17,7 +17,7 @@ import play.api.mvc.{ Request, Result, Results }
 import org.elastic4play.ErrorHandler
 
 class Renderer @Inject() (
-  errorHandler: ErrorHandler,
+    errorHandler: ErrorHandler,
     implicit val ec: ExecutionContext,
     implicit val mat: Materializer) {
 

diff --git a/app/org/elastic4play/services/AuxSrv.scala b/app/org/elastic4play/services/AuxSrv.scala
@@ -12,7 +12,7 @@ import play.api.libs.json.{ JsObject, Json }
 
 import org.elastic4play.InternalError
 import org.elastic4play.database.DBConfiguration
-import org.elastic4play.models.{ BaseEntity, ChildModelDef }
+import org.elastic4play.models.{ AttributeOption, BaseEntity, ChildModelDef }
 import org.elastic4play.models.JsonFormat.baseModelEntityWrites
 
 @Singleton
@@ -25,14 +25,24 @@ class AuxSrv @Inject() (
   import QueryDSL._
   val log = Logger(getClass)
 
-  def apply(entity: BaseEntity, nparent: Int, withStats: Boolean): Future[JsObject] = {
+  def removeUnauditedAttributes(entity: BaseEntity): JsObject = {
+    JsObject(
+      entity.attributes.fields
+        .map { case (name, value) ⇒ (name, value, entity.model.attributes.find(_.name == name)) }
+        .collect { case (name, value, Some(desc)) if !desc.options.contains(AttributeOption.unaudited) ⇒ name → value })
+  }
+  def apply(entity: BaseEntity, nparent: Int, withStats: Boolean, removeUnaudited: Boolean): Future[JsObject] = {
     val entityWithParent = entity.model match {
       case childModel: ChildModelDef[_, _, _, _] if nparent > 0 ⇒
         val (src, total) = findSrv(childModel.parentModel, ("_id" ~= entity.parentId.getOrElse(throw InternalError(s"Child entity $entity has no parent ID"))), Some("0-1"), Nil)
         src
           .mapAsync(1) { parent ⇒
-            apply(parent, nparent - 1, withStats).map { parent ⇒
-              Json.toJson(entity).as[JsObject] + (childModel.parentModel.name → parent)
+            apply(parent, nparent - 1, withStats, removeUnaudited).map { parent ⇒
+              val entityObj = removeUnaudited match {
+                case true  ⇒ removeUnauditedAttributes(entity)
+                case false ⇒ Json.toJson(entity).as[JsObject]
+              }
+              entityObj + (childModel.parentModel.name → parent)
             }
           }
           .runWith(Sink.headOption)
@@ -51,17 +61,17 @@ class AuxSrv @Inject() (
     else entityWithParent
   }
 
-  def apply[A](entities: Source[BaseEntity, A], nparent: Int, withStats: Boolean): Source[JsObject, A] = {
-    entities.mapAsync(5) { entity ⇒ apply(entity, nparent, withStats) }
+  def apply[A](entities: Source[BaseEntity, A], nparent: Int, withStats: Boolean, removeUnaudited: Boolean): Source[JsObject, A] = {
+    entities.mapAsync(5) { entity ⇒ apply(entity, nparent, withStats, removeUnaudited) }
   }
 
-  def apply(modelName: String, entityId: String, nparent: Int, withStats: Boolean): Future[JsObject] = {
+  def apply(modelName: String, entityId: String, nparent: Int, withStats: Boolean, removeUnaudited: Boolean): Future[JsObject] = {
     if (entityId == "")
       return Future.successful(JsObject(Nil))
     modelSrv(modelName)
       .map { model ⇒
         val (src, total) = findSrv(model, ("_id" ~= entityId), Some("0-1"), Nil)
-        src.mapAsync(1) { entity ⇒ apply(entity, nparent, withStats) }
+        src.mapAsync(1) { entity ⇒ apply(entity, nparent, withStats, removeUnaudited) }
           .runWith(Sink.headOption)
           .map(_.getOrElse {
             log.warn(s"Entity $modelName $entityId not found")