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
85
86
|
# ========================================================================== #
# #
# KVMD - The main PiKVM daemon. #
# #
# Copyright (C) 2018-2021 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 Dict
from typing import AsyncGenerator
from typing import Optional
import aiohttp
import aiohttp.multipart
from . import __version__
# =====
def make_user_agent(app: str) -> str:
return f"{app}/{__version__}"
def raise_not_200(response: aiohttp.ClientResponse) -> None:
if response.status != 200:
assert response.reason is not None
response.release()
raise aiohttp.ClientResponseError(
response.request_info,
response.history,
status=response.status,
message=response.reason,
headers=response.headers,
)
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")
@contextlib.asynccontextmanager
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
|
