1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
|
# ========================================================================== #
# #
# KVMD - The main PiKVM daemon. #
# #
# Copyright (C) 2018-2024 Maxim Devaev <[email protected]> #
# #
# This program is free software: you can redistribute it and/or modify #
# it under the terms of the GNU General Public License as published by #
# the Free Software Foundation, either version 3 of the License, or #
# (at your option) any later version. #
# #
# This program is distributed in the hope that it will be useful, #
# but WITHOUT ANY WARRANTY; without even the implied warranty of #
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the #
# GNU General Public License for more details. #
# #
# You should have received a copy of the GNU General Public License #
# along with this program. If not, see <https://www.gnu.org/licenses/>. #
# #
# ========================================================================== #
import os
import contextlib
from typing import AsyncGenerator
import aiohttp
import aiohttp.multipart
from . import __version__
# =====
def make_user_agent(app: str) -> str:
return f"{app}/{__version__}"
def raise_not_200(resp: aiohttp.ClientResponse) -> None:
if resp.status != 200:
assert resp.reason is not None
resp.release()
raise aiohttp.ClientResponseError(
resp.request_info,
resp.history,
status=resp.status,
message=resp.reason,
headers=resp.headers,
)
def get_filename(resp: aiohttp.ClientResponse) -> str:
try:
disp = resp.headers["Content-Disposition"]
parsed = aiohttp.multipart.parse_content_disposition(disp)
return str(parsed[1]["filename"])
except Exception:
try:
return os.path.basename(resp.url.path)
except Exception:
raise aiohttp.ClientError("Can't determine filename")
@contextlib.asynccontextmanager
async def download(
url: str,
verify: bool=True,
timeout: float=10.0,
read_timeout: (float | None)=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 resp: # type: ignore
raise_not_200(resp)
yield resp
|
