#170 Replace AIV into FIV in alert processing

TheHive-Project · May 11, 2017 · 225ba55 · 225ba55
1 parent c6b4cf5
commit 225ba55
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 22 deletions.
diff --git a/thehive-backend/app/services/AlertSrv.scala b/thehive-backend/app/services/AlertSrv.scala
@@ -14,7 +14,7 @@ import play.api.{ Configuration, Logger }
 import play.api.libs.json._
 
 import scala.concurrent.{ ExecutionContext, Future }
-import scala.util.{ Success, Try }
+import scala.util.{ Failure, Try }
 
 trait AlertTransformer {
   def createCase(alert: Alert)(implicit authContext: AuthContext): Future[Case]
@@ -142,42 +142,51 @@ class AlertSrv(
                 .set("status", CaseStatus.Open.toString))
                 .flatMap { caze ⇒ setCase(alert, caze).map(_ ⇒ caze) }
                 .flatMap { caze ⇒
-                  Future
-                    .traverse(alert.artifacts()) { artifact ⇒
+                  val artifactsFields = alert.artifacts()
+                    .map { artifact ⇒
                       val tags = (artifact \ "tags").asOpt[Seq[JsString]].getOrElse(Nil) :+ JsString("src:" + alert.tpe())
                       val message = (artifact \ "message").asOpt[JsString].getOrElse(JsString(""))
                       val artifactFields = Fields(artifact +
                         ("tags" → JsArray(tags)) +
                         ("message" → message))
                       if (artifactFields.getString("dataType").contains("file")) {
                         artifactFields.getString("data")
-                          .flatMap {
+                          .map {
                             case dataExtractor(filename, contentType, data) ⇒
                               val f = Files.createTempFile("alert-", "-attachment")
                               Files.write(f, java.util.Base64.getDecoder.decode(data))
-                              val fiv = FileInputValue(filename, f, contentType)
-                              Some(attachmentSrv
-                                .save(fiv)
-                                .map { attachment ⇒
-                                  artifactFields
-                                    .set("attachment", AttachmentInputValue(attachment))
-                                    .unset("data")
-                                }
-                                .andThen {
-                                  case _ ⇒ Files.delete(f)
-                                })
-                            case _ ⇒ None
+                              artifactFields
+                                .set("attachment", FileInputValue(filename, f, contentType))
+                                .unset("data")
+                            case data ⇒
+                              logger.warn(s"Invalid data format for file artifact: $data")
+                              artifactFields
                           }
-                          .getOrElse(Future.successful(artifactFields))
+                          .getOrElse(artifactFields)
                       }
                       else {
-                        Future.successful(artifactFields)
+                        artifactFields
                       }
                     }
-                    .flatMap { artifactsFields ⇒
-                      artifactSrv.create(caze, artifactsFields)
+
+                  val createdCase = artifactSrv.create(caze, artifactsFields)
+                    .map { r ⇒
+                      r.foreach {
+                        case Failure(e) ⇒ logger.warn("Create artifact error", e)
+                        case _          ⇒
+                      }
+                      caze
                     }
-                    .map(_ ⇒ caze)
+                  createdCase.onComplete {
+                    // remove temporary files
+                    case _ ⇒ artifactsFields
+                      .flatMap(_.get("Attachment"))
+                      .foreach {
+                        case FileInputValue(_, file, _) ⇒ Files.delete(file)
+                        case _                          ⇒
+                      }
+                  }
+                  createdCase
                 }
             }
         }

diff --git a/thehive-backend/app/services/ArtifactSrv.scala b/thehive-backend/app/services/ArtifactSrv.scala
@@ -15,6 +15,7 @@ import models.{ Artifact, ArtifactModel, ArtifactStatus, Case, CaseModel }
 import org.elastic4play.utils.{ RichFuture, RichOr }
 import models.CaseStatus
 import models.CaseResolutionStatus
+import play.api.Logger
 import play.api.libs.json.JsObject
 
 @Singleton
@@ -29,6 +30,8 @@ class ArtifactSrv @Inject() (
     fieldsSrv: FieldsSrv,
     implicit val ec: ExecutionContext) {
 
+  private[ArtifactSrv] lazy val logger = Logger(getClass)
+
   def create(caseId: String, fields: Fields)(implicit authContext: AuthContext): Future[Artifact] =
     getSrv[CaseModel, Case](caseModel, caseId)
       .flatMap { caze ⇒ create(caze, fields) }
@@ -47,7 +50,7 @@ class ArtifactSrv @Inject() (
         updatedArtifact ← updateSrv[ArtifactModel, Artifact](artifactModel, artifact.id, fields.unset("data").unset("dataType").unset("attachment").set("status", "Ok"))
       } yield updatedArtifact
       updatedArtifact.recoverWith {
-        case _ ⇒ Future.failed(CreateError(Some("CONFLICT"), "Artifact already exists", attrs))
+        case _ ⇒ Future.failed(ConflictError("Artifact already exists", attrs))
       }
     }
   }