import os import re import tempfile from pathlib import Path from urllib.parse import urlencode, unquote import aiohttp from aiohttp_socks import ProxyConnector, ProxyType from config import PROXY_URL, PROXY_ENABLED DOWNLOAD_DIR = Path("bot/data/downloads") DOWNLOAD_DIR.mkdir(parents=True, exist_ok=True) def _get_connector(): if PROXY_ENABLED and PROXY_URL: parsed = PROXY_URL.replace("socks5://", "").replace("socks5h://", "") if "@" in parsed: auth, host_port = parsed.split("@", 1) username, password = auth.split(":", 1) else: username = None password = None host_port = parsed host, port = host_port.rsplit(":", 1) port = int(port) return ProxyConnector( proxy_type=ProxyType.SOCKS5, host=host, port=port, username=username, password=password, ) return None def sanitize_filename(filename: str) -> str: filename = os.path.basename(filename).strip() filename = re.sub(r"[\\/:*?\"<>|]+", "_", filename) return filename or "downloaded_file" async def download_yandex_file(public_url: str, progress_callback=None) -> str: base_url = "https://cloud-api.yandex.net/v1/disk/public/resources/download?" final_url = base_url + urlencode({"public_key": public_url}) timeout = aiohttp.ClientTimeout(total=None, sock_connect=30, sock_read=300) connector = _get_connector() async with aiohttp.ClientSession(timeout=timeout, connector=connector) as session: async with session.get(final_url) as response: response.raise_for_status() payload = await response.json() download_url = payload["href"] async with session.get(download_url) as response: response.raise_for_status() content_disposition = response.headers.get("Content-Disposition", "") filename = download_url.split("/")[-1] if "filename*" in content_disposition: try: encoded = content_disposition.split("filename*=")[1].strip() parts = encoded.split("''", 1) if len(parts) == 2: filename = unquote(parts[1], encoding=parts[0] or "utf-8") except Exception: pass filename = sanitize_filename(filename) suffix = Path(filename).suffix or ".bin" with tempfile.NamedTemporaryFile(delete=False, suffix=suffix, dir=str(DOWNLOAD_DIR)) as tmp_file: total_size = int(response.headers.get("Content-Length", 0)) downloaded_size = 0 async for chunk in response.content.iter_chunked(1024 * 64): tmp_file.write(chunk) downloaded_size += len(chunk) if progress_callback: await progress_callback(downloaded_size, total_size) return tmp_file.name