Files
umb/bot/utils/yadisk_download.py
T
Галингер Р.С. 0f674f8832 refactor: review fixes and code improvements
- Add shared proxy module (bot/utils/proxy.py) to eliminate duplicate SOCKS5 parsing
- Fix AI client: escape HTML before Markdown→HTML conversion, unify timeouts,
  make health check optional and disabled by default, handle 429 retries
- Fix ai.py: correct forwarded message handling (aiogram 3.x forward_origin),
  pass relevant summaries via extra_system_content
- Fix dialogue.py: only respond to Astra's messages, use system context instead
  of prompt injection, answer on the limit message before phase transition
- Fix voice.py: load Whisper model in thread pool, safe WAV path generation
- Improve database.py: composite indexes, Boolean is_active, upsert file_id cache,
  add context cleanup helper
- Update weather.py and yadisk_download.py to use shared proxy connector
- Update yadisk.py: validate URL before cache clear, add download size limit,
  wrap sync file ops in to_thread
- Reuse S3 client via lru_cache
- Update setup_commands with /aiclear and /aiuser
- Update README, Dockerfile (Python 3.11), docker-compose (mount models)
- Pin dependency versions, remove unused httpx[socks]
- Add basic pytest tests for layout converter and voice normalization
2026-07-07 18:42:10 +07:00

69 lines
2.5 KiB
Python

import os
import re
import tempfile
from pathlib import Path
from urllib.parse import urlencode, unquote
import aiohttp
from bot.utils.proxy import get_proxy_connector
DOWNLOAD_DIR = Path("bot/data/downloads")
DOWNLOAD_DIR.mkdir(parents=True, exist_ok=True)
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_proxy_connector()
try:
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
finally:
if connector:
await connector.close()