#1842 Replace forbidden characters instead of refuse the download

TheHive-Project · Mar 18, 2021 · 6f1b9cc · 6f1b9cc
1 parent 08062a9
commit 6f1b9cc
Showing 1 changed file with 55 additions and 59 deletions.
diff --git a/thehive/app/org/thp/thehive/controllers/v0/AttachmentCtrl.scala b/thehive/app/org/thp/thehive/controllers/v0/AttachmentCtrl.scala
@@ -33,72 +33,68 @@ class AttachmentCtrl @Inject() (
   def download(id: String, name: Option[String]): Action[AnyContent] =
     entrypoint("download attachment")
       .authRoTransaction(db) { implicit authContext => implicit graph =>
-        if (name.getOrElse("").intersect(forbiddenChar).nonEmpty)
-          Success(Results.BadRequest("File name is invalid"))
-        else
-          attachmentSrv
-            .get(EntityIdOrName(id))
-            .visible
-            .getOrFail("Attachment")
-            .filter(attachmentSrv.exists)
-            .map { attachment =>
-              Result(
-                header = ResponseHeader(
-                  200,
-                  Map(
-                    "Content-Disposition"       -> s"""attachment; ${HttpHeaderParameterEncoding.encode("filename", name.getOrElse(id))}""",
-                    "Content-Transfer-Encoding" -> "binary"
-                  )
-                ),
-                body = HttpEntity.Streamed(attachmentSrv.source(attachment), None, None)
-              )
-            }
-            .recoverWith {
-              case _: NoSuchElementException => Failure(NotFoundError(s"Attachment $id not found"))
-            }
+        val filename = name.getOrElse(id).map(c => if (forbiddenChar.contains(c)) '_' else c)
+        attachmentSrv
+          .get(EntityIdOrName(id))
+          .visible
+          .getOrFail("Attachment")
+          .filter(attachmentSrv.exists)
+          .map { attachment =>
+            Result(
+              header = ResponseHeader(
+                200,
+                Map(
+                  "Content-Disposition"       -> s"""attachment; ${HttpHeaderParameterEncoding.encode("filename", filename)}""",
+                  "Content-Transfer-Encoding" -> "binary"
+                )
+              ),
+              body = HttpEntity.Streamed(attachmentSrv.source(attachment), None, None)
+            )
+          }
+          .recoverWith {
+            case _: NoSuchElementException => Failure(NotFoundError(s"Attachment $id not found"))
+          }
       }
 
   def downloadZip(id: String, name: Option[String]): Action[AnyContent] =
     entrypoint("download attachment")
       .authRoTransaction(db) { implicit authContext => implicit graph =>
-        if (name.getOrElse("").intersect(forbiddenChar).nonEmpty)
-          Success(Results.BadRequest("File name is invalid"))
-        else
-          attachmentSrv
-            .get(EntityIdOrName(id))
-            .visible
-            .getOrFail("Attachment")
-            .filter(attachmentSrv.exists)
-            .flatMap { attachment =>
-              Try {
-                val f = Files.createTempFile("downloadzip-", id)
-                Files.delete(f)
-                val zipFile   = new ZipFile(f.toFile, password.toCharArray)
-                val zipParams = new ZipParameters
-                zipParams.setCompressionLevel(CompressionLevel.FASTEST)
-                zipParams.setEncryptFiles(true)
-                zipParams.setEncryptionMethod(EncryptionMethod.ZIP_STANDARD)
-                zipParams.setFileNameInZip(name.getOrElse(id))
-                //      zipParams.setSourceExternalStream(true)
-                zipFile.addStream(attachmentSrv.stream(attachment), zipParams)
+        val filename = name.getOrElse(id).map(c => if (forbiddenChar.contains(c)) '_' else c)
+        attachmentSrv
+          .get(EntityIdOrName(id))
+          .visible
+          .getOrFail("Attachment")
+          .filter(attachmentSrv.exists)
+          .flatMap { attachment =>
+            Try {
+              val f = Files.createTempFile("downloadzip-", id)
+              Files.delete(f)
+              val zipFile   = new ZipFile(f.toFile, password.toCharArray)
+              val zipParams = new ZipParameters
+              zipParams.setCompressionLevel(CompressionLevel.FASTEST)
+              zipParams.setEncryptFiles(true)
+              zipParams.setEncryptionMethod(EncryptionMethod.ZIP_STANDARD)
+              zipParams.setFileNameInZip(filename)
+              //      zipParams.setSourceExternalStream(true)
+              zipFile.addStream(attachmentSrv.stream(attachment), zipParams)
 
-                Result(
-                  header = ResponseHeader(
-                    200,
-                    Map(
-                      "Content-Disposition"       -> s"""attachment; filename="${name.getOrElse(id)}.zip"""",
-                      "Content-Type"              -> "application/zip",
-                      "Content-Transfer-Encoding" -> "binary",
-                      "Content-Length"            -> Files.size(f).toString
-                    )
-                  ),
-                  body = HttpEntity.Streamed(FileIO.fromPath(f), Some(Files.size(f)), Some("application/zip"))
-                ) // FIXME remove temporary file (but when ?)
-              }
-            }
-            .recoverWith {
-              case _: NoSuchElementException => Failure(NotFoundError(s"Attachment $id not found"))
+              Result(
+                header = ResponseHeader(
+                  200,
+                  Map(
+                    "Content-Disposition"       -> s"""attachment; filename="$filename.zip"""",
+                    "Content-Type"              -> "application/zip",
+                    "Content-Transfer-Encoding" -> "binary",
+                    "Content-Length"            -> Files.size(f).toString
+                  )
+                ),
+                body = HttpEntity.Streamed(FileIO.fromPath(f), Some(Files.size(f)), Some("application/zip"))
+              ) // FIXME remove temporary file (but when ?)
             }
+          }
+          .recoverWith {
+            case _: NoSuchElementException => Failure(NotFoundError(s"Attachment $id not found"))
+          }
       }
 
   def password: String = passwordConfig.get