| import os |
| import re |
| import httpx |
| import uuid |
| from config import DOWNLOAD_DIR, MAX_FILE_SIZE |
|
|
| os.makedirs(DOWNLOAD_DIR, exist_ok=True) |
|
|
|
|
| async def download_file(url: str): |
| async with httpx.AsyncClient(timeout=300, follow_redirects=True) as client: |
| r = await client.head(url) |
| content_length = int(r.headers.get("content-length", 0)) |
| content_type = r.headers.get("content-type", "application/octet-stream") |
|
|
| r = await client.get(url) |
| if r.status_code != 200: |
| return None, f"Failed to download (HTTP {r.status_code})", None, None |
|
|
| filename = _extract_filename(r.headers, url, content_type) |
| if not filename: |
| ext = _ext_from_url(url, content_type) |
| filename = f"{uuid.uuid4()}{ext}" |
|
|
| filepath = os.path.join(DOWNLOAD_DIR, filename) |
|
|
| with open(filepath, "wb") as f: |
| f.write(r.content) |
|
|
| return filepath, None, content_type, len(r.content) |
|
|
|
|
| def _extract_filename(headers, url: str, content_type: str) -> str: |
| cd = headers.get("Content-Disposition", "") |
| match = re.search(r'filename\*?="?([^";]+)"?', cd) |
| if match: |
| name = match.group(1).strip() |
| if name: |
| return name |
|
|
| from urllib.parse import urlparse, unquote |
| path = urlparse(url).path |
| if "/" in path: |
| name = path.rsplit("/", 1)[-1] |
| if name and "." in name: |
| return unquote(name) |
|
|
| ext = _ext_from_url(url, content_type) |
| if ext: |
| return f"download{ext}" |
| return "" |
|
|
|
|
| def _ext_from_url(url: str, content_type: str) -> str: |
| match = re.search(r"\.([a-zA-Z0-9]+)(?:\?|$)", url) |
| if match: |
| return f".{match.group(1)}" |
| mapping = { |
| "image/jpeg": ".jpg", |
| "image/png": ".png", |
| "image/gif": ".gif", |
| "application/pdf": ".pdf", |
| "application/zip": ".zip", |
| "text/plain": ".txt", |
| "application/x-msdownload": ".exe", |
| "application/vnd.ms-excel": ".xls", |
| "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet": ".xlsx", |
| "application/msword": ".doc", |
| "application/vnd.openxmlformats-officedocument.wordprocessingml.document": ".docx", |
| "video/mp4": ".mp4", |
| "audio/mpeg": ".mp3", |
| } |
| return mapping.get(content_type, "") |
|
|
|
|
| def cleanup_file(filepath: str): |
| try: |
| if os.path.exists(filepath): |
| os.remove(filepath) |
| except Exception: |
| pass |
|
|
