#1697 Fix observable parsing with attachment

TheHive-Project · Dec 15, 2020 · 58a3d0a · 58a3d0a
1 parent 553736a
commit 58a3d0a
Show file tree

Hide file tree

Showing 8 changed files with 212 additions and 92 deletions.
diff --git a/ScalliGraph b/ScalliGraph
diff --git a/dto/src/main/scala/org/thp/thehive/dto/v0/Attachment.scala b/dto/src/main/scala/org/thp/thehive/dto/v0/Attachment.scala
@@ -1,6 +1,12 @@
 package org.thp.thehive.dto.v0
 
-import play.api.libs.json.{Json, OFormat}
+import play.api.libs.json.{Json, OFormat, Writes}
+
+case class InputAttachment(name: String, contentType: String, id: String)
+
+object InputAttachment {
+  implicit val writes: Writes[InputAttachment] = Json.writes[InputAttachment]
+}
 
 case class OutputAttachment(name: String, hashes: Seq[String], size: Long, contentType: String, id: String)
 

diff --git a/dto/src/main/scala/org/thp/thehive/dto/v0/Observable.scala b/dto/src/main/scala/org/thp/thehive/dto/v0/Observable.scala
@@ -1,19 +1,19 @@
 package org.thp.thehive.dto.v0
 
 import java.util.Date
-
 import org.scalactic.Accumulation._
 import org.scalactic.Good
 import org.thp.scalligraph.controllers._
 import play.api.libs.json.{JsObject, Json, OFormat, Writes}
 
 case class InputObservable(
     dataType: String,
-    @WithParser(InputObservable.fp)
+    @WithParser(InputObservable.dataParser)
     data: Seq[String] = Nil,
     message: Option[String] = None,
     startDate: Option[Date] = None,
-    attachment: Option[FFile] = None,
+    @WithParser(InputObservable.fileOrAttachmentParser)
+    attachment: Seq[Either[FFile, InputAttachment]] = Seq.empty,
     tlp: Option[Int] = None,
     tags: Set[String] = Set.empty,
     ioc: Option[Boolean] = None,
@@ -22,14 +22,31 @@ case class InputObservable(
 )
 
 object InputObservable {
+  implicit val fileOrAttachmentWrites: Writes[Either[FFile, InputAttachment]] = Writes[Either[FFile, InputAttachment]] {
+    case Left(file)        => Json.toJson(file)
+    case Right(attachment) => Json.toJson(attachment)
+  }
   implicit val writes: Writes[InputObservable] = Json.writes[InputObservable]
 
-  val fp: FieldsParser[Seq[String]] = FieldsParser[Seq[String]]("data") {
+  val dataParser: FieldsParser[Seq[String]] = FieldsParser[Seq[String]]("data") {
     case (_, FString(s)) => Good(Seq(s))
     case (_, FAny(s))    => Good(s)
     case (_, FSeq(a))    => a.validatedBy(FieldsParser.string(_))
     case (_, FUndefined) => Good(Nil)
   }
+
+  val fileOrAttachmentParser: FieldsParser[Seq[Either[FFile, InputAttachment]]] =
+    FieldsParser[FFile]
+      .map("fileOrAttachmentParser")(f => Seq(Left(f)))
+      .recover(
+        FieldsParser[InputAttachment]
+          .map("fileOrAttachmentParser")(a => Seq(Right(a)))
+          .recover(
+            FieldsParser[InputAttachment]
+              .sequence
+              .map("fileOrAttachmentParser")(as => as.map(Right(_)))
+          )
+      )
 }
 
 case class OutputObservable(

diff --git a/dto/src/main/scala/org/thp/thehive/dto/v1/Attachment.scala b/dto/src/main/scala/org/thp/thehive/dto/v1/Attachment.scala
@@ -1,6 +1,12 @@
 package org.thp.thehive.dto.v1
 
-import play.api.libs.json.{Json, OFormat}
+import play.api.libs.json.{Json, OFormat, Writes}
+
+case class InputAttachment(name: String, contentType: String, id: String)
+
+object InputAttachment {
+  implicit val writes: Writes[InputAttachment] = Json.writes[InputAttachment]
+}
 
 case class OutputAttachment(name: String, hashes: Seq[String], size: Long, contentType: String, id: String)
 

diff --git a/dto/src/main/scala/org/thp/thehive/dto/v1/Observable.scala b/dto/src/main/scala/org/thp/thehive/dto/v1/Observable.scala
@@ -1,19 +1,19 @@
 package org.thp.thehive.dto.v1
 
 import java.util.Date
-
 import org.scalactic.Accumulation._
 import org.scalactic.Good
 import org.thp.scalligraph.controllers._
 import play.api.libs.json.{JsObject, Json, OFormat, Writes}
 
 case class InputObservable(
     dataType: String,
-    @WithParser(InputObservable.fp)
+    @WithParser(InputObservable.dataParser)
     data: Seq[String] = Nil,
     message: Option[String] = None,
     startDate: Option[Date] = None,
-    attachment: Option[FFile] = None,
+    @WithParser(InputObservable.fileOrAttachmentParser)
+    attachment: Seq[Either[FFile, InputAttachment]] = Seq.empty,
     tlp: Option[Int] = None,
     tags: Set[String] = Set.empty,
     ioc: Option[Boolean] = None,
@@ -22,14 +22,32 @@ case class InputObservable(
 )
 
 object InputObservable {
+  implicit val fileOrAttachmentWrites: Writes[Either[FFile, InputAttachment]] = Writes[Either[FFile, InputAttachment]] {
+    case Left(file)        => Json.toJson(file)
+    case Right(attachment) => Json.toJson(attachment)
+  }
+
   implicit val writes: Writes[InputObservable] = Json.writes[InputObservable]
 
-  val fp: FieldsParser[Seq[String]] = FieldsParser[Seq[String]]("data") {
+  val dataParser: FieldsParser[Seq[String]] = FieldsParser[Seq[String]]("data") {
     case (_, FString(s)) => Good(Seq(s))
     case (_, FAny(s))    => Good(s)
     case (_, FSeq(a))    => a.validatedBy(FieldsParser.string(_))
     case (_, FUndefined) => Good(Nil)
   }
+
+  val fileOrAttachmentParser: FieldsParser[Seq[Either[FFile, InputAttachment]]] =
+    FieldsParser[FFile]
+      .map("fileOrAttachmentParser")(f => Seq(Left(f)))
+      .recover(
+        FieldsParser[InputAttachment]
+          .map("fileOrAttachmentParser")(a => Seq(Right(a)))
+          .recover(
+            FieldsParser[InputAttachment]
+              .sequence
+              .map("fileOrAttachmentParser")(as => as.map(Right(_)))
+          )
+      )
 }
 
 case class OutputObservable(

diff --git a/thehive/app/org/thp/thehive/controllers/v0/ObservableCtrl.scala b/thehive/app/org/thp/thehive/controllers/v0/ObservableCtrl.scala
@@ -1,20 +1,16 @@
 package org.thp.thehive.controllers.v0
 
-import java.io.FilterInputStream
-import java.nio.file.Files
-
-import javax.inject.{Inject, Named, Singleton}
 import net.lingala.zip4j.ZipFile
 import net.lingala.zip4j.model.FileHeader
 import org.thp.scalligraph._
 import org.thp.scalligraph.auth.AuthContext
 import org.thp.scalligraph.controllers._
-import org.thp.scalligraph.models.{Database, UMapping}
+import org.thp.scalligraph.models.{Database, Entity, UMapping}
 import org.thp.scalligraph.query._
 import org.thp.scalligraph.traversal.TraversalOps._
 import org.thp.scalligraph.traversal.{Converter, IteratorOutput, Traversal}
 import org.thp.thehive.controllers.v0.Conversion._
-import org.thp.thehive.dto.v0.InputObservable
+import org.thp.thehive.dto.v0.{InputAttachment, InputObservable}
 import org.thp.thehive.models._
 import org.thp.thehive.services.CaseOps._
 import org.thp.thehive.services.ObservableOps._
@@ -27,8 +23,11 @@ import play.api.libs.Files.DefaultTemporaryFileCreator
 import play.api.libs.json.{JsArray, JsObject, JsValue, Json}
 import play.api.mvc.{Action, AnyContent, Results}
 
+import java.io.FilterInputStream
+import java.nio.file.Files
+import javax.inject.{Inject, Named, Singleton}
 import scala.collection.JavaConverters._
-import scala.util.Success
+import scala.util.{Failure, Success}
 
 @Singleton
 class ObservableCtrl @Inject() (
@@ -38,6 +37,7 @@ class ObservableCtrl @Inject() (
     observableSrv: ObservableSrv,
     observableTypeSrv: ObservableTypeSrv,
     caseSrv: CaseSrv,
+    attachmentSrv: AttachmentSrv,
     errorHandler: ErrorHandler,
     @Named("v0") override val queryExecutor: QueryExecutor,
     override val publicData: PublicObservable,
@@ -68,48 +68,60 @@ class ObservableCtrl @Inject() (
           }
           .map {
             case (case0, observableType) =>
-              val initialSuccessesAndFailures: (Seq[JsValue], Seq[JsValue]) =
-                inputAttachObs.foldLeft[(Seq[JsValue], Seq[JsValue])](Nil -> Nil) {
-                  case ((successes, failures), inputObservable) =>
-                    inputObservable.attachment.fold((successes, failures)) { attachment =>
-                      db
-                        .tryTransaction { implicit graph =>
-                          observableSrv
-                            .create(inputObservable.toObservable, observableType, attachment, inputObservable.tags, Nil)
-                            .flatMap(o => caseSrv.addObservable(case0, o).map(_ => o.toJson))
-                        }
-                        .fold(
-                          e =>
-                            successes -> (failures :+ errorHandler.toErrorResult(e)._2 ++ Json
-                              .obj(
-                                "object" -> Json
-                                  .obj("data" -> s"file:${attachment.filename}", "attachment" -> Json.obj("name" -> attachment.filename))
-                              )),
-                          s => (successes :+ s) -> failures
-                        )
-                    }
+              val (successes, failures) = inputAttachObs
+                .flatMap { obs =>
+                  obs.attachment.map(createAttachmentObservable(case0, obs, observableType, _)) ++
+                    obs.data.map(createSimpleObservable(case0, obs, observableType, _))
                 }
-
-              val (successes, failures) = inputObservable
-                .data
-                .foldLeft(initialSuccessesAndFailures) {
-                  case ((successes, failures), data) =>
-                    db
-                      .tryTransaction { implicit graph =>
-                        observableSrv
-                          .create(inputObservable.toObservable, observableType, data, inputObservable.tags, Nil)
-                          .flatMap(o => caseSrv.addObservable(case0, o).map(_ => o.toJson))
-                      }
-                      .fold(
-                        failure => (successes, failures :+ errorHandler.toErrorResult(failure)._2 ++ Json.obj("object" -> Json.obj("data" -> data))),
-                        success => (successes :+ success, failures)
-                      )
+                .foldLeft[(Seq[JsValue], Seq[JsValue])]((Nil, Nil)) {
+                  case ((s, f), Right(o)) => (s :+ o, f)
+                  case ((s, f), Left(o))  => (s, f :+ o)
                 }
               if (failures.isEmpty) Results.Created(JsArray(successes))
               else Results.MultiStatus(Json.obj("success" -> successes, "failure" -> failures))
           }
       }
 
+  def createSimpleObservable(
+      `case`: Case with Entity,
+      inputObservable: InputObservable,
+      observableType: ObservableType with Entity,
+      data: String
+  )(implicit authContext: AuthContext): Either[JsValue, JsValue] =
+    db
+      .tryTransaction { implicit graph =>
+        observableSrv
+          .create(inputObservable.toObservable, observableType, data, inputObservable.tags, Nil)
+          .flatMap(o => caseSrv.addObservable(`case`, o).map(_ => o))
+      } match {
+      case Success(o)     => Right(o.toJson)
+      case Failure(error) => Left(errorHandler.toErrorResult(error)._2 ++ Json.obj("object" -> Json.obj("data" -> data)))
+    }
+
+  def createAttachmentObservable(
+      `case`: Case with Entity,
+      inputObservable: InputObservable,
+      observableType: ObservableType with Entity,
+      fileOrAttachment: Either[FFile, InputAttachment]
+  )(implicit authContext: AuthContext): Either[JsValue, JsValue] =
+    db
+      .tryTransaction { implicit graph =>
+        val observable = fileOrAttachment match {
+          case Left(file) => observableSrv.create(inputObservable.toObservable, observableType, file, inputObservable.tags, Nil)
+          case Right(attachment) =>
+            for {
+              attach <- attachmentSrv.duplicate(attachment.name, attachment.contentType, attachment.id)
+              obs    <- observableSrv.create(inputObservable.toObservable, observableType, attach, inputObservable.tags, Nil)
+            } yield obs
+        }
+        observable.flatMap(o => caseSrv.addObservable(`case`, o).map(_ => o))
+      } match {
+      case Success(o) => Right(o.toJson)
+      case _ =>
+        val filename = fileOrAttachment.fold(_.filename, _.name)
+        Left(Json.obj("object" -> Json.obj("data" -> s"file:$filename", "attachment" -> Json.obj("name" -> filename))))
+    }
+
   def get(observableId: String): Action[AnyContent] =
     entrypoint("get observable")
       .authRoTransaction(db) { implicit request => implicit graph =>
@@ -214,8 +226,8 @@ class ObservableCtrl @Inject() (
     }
   }
 
-  private def getZipFiles(observable: InputObservable, zipPassword: Option[String])(implicit authContext: AuthContext): Seq[InputObservable] =
-    observable.attachment.toSeq.flatMap { attachment =>
+  private def getZipFiles(observable: InputObservable, zipPassword: Option[String]): Seq[InputObservable] =
+    observable.attachment.flatMap(_.swap.toSeq).flatMap { attachment =>
       val zipFile                = new ZipFile(attachment.filepath.toFile)
       val files: Seq[FileHeader] = zipFile.getFileHeaders.asScala.asInstanceOf[Seq[FileHeader]]
 
@@ -225,7 +237,7 @@ class ObservableCtrl @Inject() (
       files
         .filterNot(_.isDirectory)
         .flatMap(extractAndCheckSize(zipFile, _))
-        .map(ffile => observable.copy(attachment = Some(ffile)))
+        .map(ffile => observable.copy(attachment = Seq(Left(ffile))))
     }
 }
+18 −13		core/src/main/scala/org/thp/scalligraph/controllers/FieldsParser.scala
+22 −0		core/src/main/scala/org/thp/scalligraph/services/StorageSrv.scala