/msd/write_remote handle

1 files changed, 53 insertions, 0 deletions

diff --git a/kvmd/htclient.py b/kvmd/htclient.py
index 368bf668..1927ede2 100644
--- a/kvmd/htclient.py
+++ b/kvmd/htclient.py
@@ -20,7 +20,15 @@
 # ========================================================================== #
 
 
+import os
+import contextlib
+
+from typing import Dict
+from typing import AsyncGenerator
+from typing import Optional
+
 import aiohttp
+import aiohttp.multipart
 
 from . import __version__
 
@@ -41,3 +49,48 @@ def raise_not_200(response: aiohttp.ClientResponse) -> None:
             message=response.reason,
             headers=response.headers,
         )
+
+
+def get_content_length(response: aiohttp.ClientResponse) -> int:
+    try:
+        value = int(response.headers["Content-Length"])
+    except Exception:
+        raise aiohttp.ClientError("Empty or invalid Content-Length")
+    if value < 0:
+        raise aiohttp.ClientError("Negative Content-Length")
+    return value
+
+
+def get_filename(response: aiohttp.ClientResponse) -> str:
+    try:
+        disp = response.headers["Content-Disposition"]
+        parsed = aiohttp.multipart.parse_content_disposition(disp)
+        return str(parsed[1]["filename"])
+    except Exception:
+        try:
+            return os.path.basename(response.url.path)
+        except Exception:
+            raise aiohttp.ClientError("Can't determine filename")
+
+
[email protected]
+async def download(
+    url: str,
+    verify: bool=True,
+    timeout: float=10.0,
+    read_timeout: Optional[float]=None,
+    app: str="KVMD",
+) -> AsyncGenerator[aiohttp.ClientResponse, None]:
+
+    kwargs: Dict = {
+        "headers": {"User-Agent": make_user_agent(app)},
+        "timeout": aiohttp.ClientTimeout(
+            connect=timeout,
+            sock_connect=timeout,
+            sock_read=(read_timeout if read_timeout is not None else timeout),
+        ),
+    }
+    async with aiohttp.ClientSession(**kwargs) as session:
+        async with session.get(url, verify_ssl=verify) as response:
+            raise_not_200(response)
+            yield response