#53 Add API to import a bundle of report templates

project/Dependencies.scala

@@ -29,7 +29,7 @@ object Dependencies {
     val reflections = "org.reflections" % "reflections" % "0.9.10"
     val zip4j = "net.lingala.zip4j" % "zip4j" % "1.3.2"
     val akkaTest = "com.typesafe.akka" %% "akka-stream-testkit" % "2.4.4"
-    val elastic4play = "org.cert-bdf" %% "elastic4play" % "1.1.1"
+    val elastic4play = "org.cert-bdf" %% "elastic4play" % "1.1.2-SNAPSHOT"
 
     object Elastic4s {
       private val version = "2.3.0"

thehive-cortex/app/connectors/cortex/models/JsonFormat.scala

@@ -7,6 +7,7 @@ import play.api.libs.json.{ OFormat, OWrites, Reads, Writes }
 import play.api.libs.json.Json.toJsFieldJsValueWrapper
 
 import org.elastic4play.models.JsonFormat.enumFormat
+import java.util.Date
 
 object JsonFormat {
   val analyzerWrites = Writes[Analyzer](analyzer ⇒ Json.obj(
@@ -46,5 +47,15 @@ object JsonFormat {
 
   implicit val artifactFormat = OFormat(artifactReads, artifactWrites)
   implicit val jobStatusFormat = enumFormat(JobStatus)
-  implicit val cortexJobFormat = Json.format[CortexJob]
+  val cortexJobReads = Reads[CortexJob](json ⇒
+    for {
+      id ← (json \ "id").validate[String]
+      analyzerId ← (json \ "analyzerId").validate[String]
+      artifact ← (json \ "artifact").validate[CortexArtifact]
+      date ← (json \ "date").validate[Date]
+      status ← (json \ "status").validate[JobStatus.Type]
+    } yield CortexJob(id, analyzerId, artifact, date, status, Nil))
+  val cortexJobWrites = Json.writes[CortexJob]
+  implicit val cortexJobFormat = Format(cortexJobReads, cortexJobWrites)
+  implicit val reportTypeFormat = enumFormat(ReportType)
 }

thehive-cortex/app/connectors/cortex/models/ReportTemplate.scala

@@ -1,17 +1,38 @@
-package models
+package connectors.cortex.models
 
 import javax.inject.{ Inject, Singleton }
 
 import play.api.libs.json.JsObject
 
-import org.elastic4play.models.{ AttributeDef, AttributeFormat ⇒ F, EntityDef, ModelDef }
+import org.elastic4play.models.{ AttributeDef, AttributeFormat ⇒ F, AttributeOption ⇒ O, EntityDef, ModelDef }
+import org.elastic4play.BadRequestError
+import org.elastic4play.models.BaseEntity
+import play.api.libs.json.JsString
+import scala.concurrent.Future
+import org.elastic4play.models.HiveEnumeration
+
+import JsonFormat._
+
+object ReportType extends Enumeration with HiveEnumeration {
+  type Type = Value
+  val short, long = Value
+}
 
 trait ReportTemplateAttributes { _: AttributeDef ⇒
+  val reportTemplateId = attribute("_id", F.stringFmt, "Report template id", O.model)
   val content = attribute("content", F.textFmt, "Content of the template")
-  val falvor = attribute("flavor", F.stringFmt, "Flavor of the report (short or long)")
-  val analyzers = multiAttribute("analyzers", F.stringFmt, "Id of analyzers")
+  val flavor = attribute("flavor", F.enumFmt(ReportType), "Flavor of the report (short or long)")
+  val analyzers = attribute("analyzers", F.stringFmt, "Id of analyzers")
 }
 
 @Singleton
-class ReportTemplateModel @Inject() extends ModelDef[ReportTemplateModel, ReportTemplate]("reportTemplate") with ReportTemplateAttributes
+class ReportTemplateModel @Inject() extends ModelDef[ReportTemplateModel, ReportTemplate]("reportTemplate") with ReportTemplateAttributes {
+  override def creationHook(parent: Option[BaseEntity], attrs: JsObject) = {
+    val maybeId = for {
+      analyzers ← (attrs \ "analyzers").asOpt[String]
+      flavor ← (attrs \ "flavor").asOpt[String]
+    } yield analyzers + "_" + flavor
+    Future.successful(maybeId.fold(attrs)(id ⇒ attrs + ("_id" → JsString(id))))
+  }
+}
 class ReportTemplate(model: ReportTemplateModel, attributes: JsObject) extends EntityDef[ReportTemplateModel, ReportTemplate](model, attributes) with ReportTemplateAttributes

thehive-cortex/app/connectors/cortex/services/CortexSrv.scala

@@ -35,7 +35,8 @@ object CortexConfig {
       val url = configuration.getString("url").getOrElse(sys.error("url is missing")).replaceFirst("/*$", "")
       val key = configuration.getString("key").getOrElse(sys.error("key is missing"))
       Some(new CortexClient(name, url, key))
-    } catch {
+    }
+    catch {
       case NonFatal(_) ⇒ None
     }
   }
@@ -114,12 +115,12 @@ class CortexSrv @Inject() (
   def getAnalyzer(analyzerId: String): Future[Analyzer] = {
     Future
       .traverse(cortexConfig.instances) {
-        case (name, cortex) => cortex.getAnalyzer(analyzerId).map(Some(_)).fallbackTo(Future.successful(None))
+        case (name, cortex) ⇒ cortex.getAnalyzer(analyzerId).map(Some(_)).fallbackTo(Future.successful(None))
       }
-      .map { analyzers =>
+      .map { analyzers ⇒
         analyzers.foldLeft[Option[Analyzer]](None) {
-          case (Some(analyzer1), Some(analyzer2)) => Some(analyzer1.copy(cortexIds = analyzer1.cortexIds ++ analyzer2.cortexIds))
-          case (maybeAnalyzer1, maybeAnalyzer2)   => maybeAnalyzer1 orElse maybeAnalyzer2
+          case (Some(analyzer1), Some(analyzer2)) ⇒ Some(analyzer1.copy(cortexIds = analyzer1.cortexIds ++ analyzer2.cortexIds))
+          case (maybeAnalyzer1, maybeAnalyzer2)   ⇒ maybeAnalyzer1 orElse maybeAnalyzer2
         }
           .getOrElse(throw NotFoundError(s"Analyzer $analyzerId not found"))
       }

thehive-cortex/app/connectors/cortex/services/ReportTemplaceSrv.scala

@@ -1,4 +1,4 @@
-package services
+package connectors.cortex.services
 
 import javax.inject.{ Inject, Singleton }
 
@@ -14,7 +14,8 @@ import play.api.libs.json.JsObject
 import org.elastic4play.controllers.Fields
 import org.elastic4play.services.{ Agg, AuthContext, CreateSrv, DeleteSrv, FindSrv, GetSrv, QueryDef, UpdateSrv }
 
-import models.{ ReportTemplate, ReportTemplateModel }
+import connectors.cortex.models.{ ReportTemplate, ReportTemplateModel }
+import services.ArtifactSrv
 
 @Singleton
 class ReportTemplateSrv @Inject() (

thehive-cortex/app/controllers/CortextCtrl.scala

@@ -43,6 +43,7 @@ class CortextCtrl @Inject() (
     case PATCH(p"/report/template/$caseTemplateId<[^/]*>") ⇒ reportTemplateCtrl.update(caseTemplateId)
     case DELETE(p"/report/template/$caseTemplateId<[^/]*>") ⇒ reportTemplateCtrl.delete(caseTemplateId)
     case GET(p"/report/template/content/$analyzerId<[^/]*>/$flavor<[^/]*>") ⇒ reportTemplateCtrl.getContent(analyzerId, flavor)
+    case POST(p"/report/template/_import") ⇒ reportTemplateCtrl.importTemplatePackage
     case r ⇒ throw NotFoundError(s"${r.uri} not found")
   }
 

thehive-cortex/app/controllers/ReportTemplate.scala

@@ -2,22 +2,29 @@ package connectors.cortex.controllers
 
 import javax.inject.{ Inject, Singleton }
 
-import scala.concurrent.ExecutionContext
+import scala.collection.JavaConversions.asScalaBuffer
+import scala.concurrent.{ ExecutionContext, Future }
+import scala.io.Source
+import scala.util.control.NonFatal
 
+import akka.stream.Materializer
+import akka.stream.scaladsl.Sink
+
+import play.api.Logger
 import play.api.http.Status
+import play.api.libs.json.{ JsBoolean, JsObject }
 import play.api.mvc.Controller
 
-import org.elastic4play.Timed
-import org.elastic4play.controllers.{ Authenticated, FieldsBodyParser, Renderer }
+import org.elastic4play.{ BadRequestError, Timed }
+import org.elastic4play.controllers.{ Authenticated, Fields, FieldsBodyParser, FileInputValue, Renderer }
 import org.elastic4play.models.JsonFormat.baseModelEntityWrites
 import org.elastic4play.services.{ QueryDSL, QueryDef, Role }
 import org.elastic4play.services.AuxSrv
 import org.elastic4play.services.JsonFormat.queryReads
 
-import services.ReportTemplateSrv
-import play.api.Logger
-import akka.stream.scaladsl.Sink
-import akka.stream.Materializer
+import connectors.cortex.services.ReportTemplateSrv
+import net.lingala.zip4j.core.ZipFile
+import net.lingala.zip4j.model.FileHeader
 
 @Singleton
 class ReportTemplateCtrl @Inject() (
@@ -82,4 +89,41 @@ class ReportTemplateCtrl @Inject() (
     val reportTemplatesWithStats = auxSrv(reportTemplates, nparent, withStats)
     renderer.toOutput(OK, reportTemplatesWithStats, total)
   }
+
+  @Timed
+  def importTemplatePackage = authenticated(Role.write).async(fieldsBodyParser) { implicit request ⇒
+    val zipFile = request.body.get("templates") match {
+      case Some(FileInputValue(name, filepath, contentType)) ⇒ new ZipFile(filepath.toFile)
+      case _                                                 ⇒ throw BadRequestError("")
+    }
+    val importedReportTemplates: Seq[Future[(String, JsBoolean)]] = zipFile.getFileHeaders.toSeq.filter(_ != null).collect {
+      case fileHeader: FileHeader if !fileHeader.isDirectory ⇒
+        val Array(analyzerId, flavor) = (fileHeader.getFileName + "/").split("/", 2)
+        val inputStream = zipFile.getInputStream(fileHeader)
+        val content = Source.fromInputStream(inputStream).mkString
+        inputStream.close()
+
+        val reportTemplateFields = Fields.empty
+          .set("flavor", flavor)
+          .set("analyzers", analyzerId)
+          .set("content", content)
+        reportTemplateSrv.create(reportTemplateFields)
+          .recoverWith { // if creation fails, try to update
+            case NonFatal(_) ⇒
+              val reportTemplateId = analyzerId + "_" + flavor
+              reportTemplateSrv.update(reportTemplateId, Fields.empty.set("content", content))
+          }
+          .map(_.id → JsBoolean(true))
+          .recoverWith {
+            case NonFatal(e) ⇒
+              logger.error(s"The import of the report template $analyzerId ($flavor) has failed", e)
+              val reportTemplateId = analyzerId + "_" + flavor
+              Future.successful(reportTemplateId → JsBoolean(false))
+          }
+    }
+
+    Future.sequence(importedReportTemplates).map { result ⇒
+      renderer.toOutput(OK, JsObject(result))
+    }
+  }
 }

thehive-cortex/build.sbt

@@ -2,7 +2,8 @@ import Dependencies._
 
 libraryDependencies ++= Seq(
   Library.Play.ws,
-  Library.elastic4play
+  Library.elastic4play,
+  Library.zip4j
 )
 
 enablePlugins(PlayScala)