From 698d47a4bcfb4ac4f802edce13630e5a2a71e766 Mon Sep 17 00:00:00 2001
From: Bulat Kurbanov <kurbanovbul@gmail.com>
Date: Thu, 19 Jan 2023 17:40:14 +0100
Subject: [PATCH] Add x-filename-b64 header

---
 src/app/services/downloader.py |  6 +++---
 src/app/views.py               | 14 ++++++++------
 2 files changed, 11 insertions(+), 9 deletions(-)

diff --git a/src/app/services/downloader.py b/src/app/services/downloader.py
index e4af733..b9d4179 100644
--- a/src/app/services/downloader.py
+++ b/src/app/services/downloader.py
@@ -1,3 +1,4 @@
+from base64 import b64decode
 from typing import Optional
 
 import httpx
@@ -28,8 +29,7 @@ async def download(
         await client.aclose()
         return None
 
-    content_disposition = response.headers["Content-Disposition"]
-    name = content_disposition.replace("attachment; filename=", "")
+    name = b64decode(response.headers["x-filename-b64"]).decode()
 
     return response, client, name
 
@@ -47,4 +47,4 @@ async def get_filename(book_id: int, file_type: str) -> Optional[str]:
         if response.status_code != 200:
             return None
 
-        return response.text.encode("ascii", "ignore").decode("ascii")
+        return response.text
diff --git a/src/app/views.py b/src/app/views.py
index 93f635b..5f6d5e4 100644
--- a/src/app/views.py
+++ b/src/app/views.py
@@ -1,5 +1,5 @@
 import asyncio
-import base64
+from base64 import b64encode
 
 from arq.connections import ArqRedis
 from fastapi import APIRouter, Depends, HTTPException, Request, status
@@ -51,7 +51,7 @@ async def download_cached_file(request: Request, object_id: int, object_type: st
     if not cached_file:
         raise HTTPException(status_code=status.HTTP_404_NOT_FOUND)
 
-    cache_data = cached_file.data
+    cache_data: dict = cached_file.data  # type: ignore
 
     data, filename, book = await asyncio.gather(
         download_file_from_cache(cache_data["chat_id"], cache_data["message_id"]),
@@ -76,14 +76,16 @@ async def download_cached_file(request: Request, object_id: int, object_type: st
         await client.aclose()
 
     assert book
+    assert filename
+
+    filename_ascii = filename.encode("ascii", "ignore").decode("ascii")
 
     return StreamingResponse(
         response.aiter_bytes(),
         headers={
-            "Content-Disposition": f"attachment; filename={filename}",
-            "X-Caption-B64": base64.b64encode(get_caption(book).encode("utf-8")).decode(
-                "latin-1"
-            ),
+            "Content-Disposition": f"attachment; filename={filename_ascii}",
+            "X-Caption-B64": b64encode(get_caption(book).encode("utf-8")).decode(),
+            "X-Filename-B64": b64encode(filename.encode("utf-8")).decode(),
         },
         background=BackgroundTask(close),
     )