| 1 | """Safe HTTP download helper — no redirects, size cap, status 200 only.""" |
| 2 | |
| 3 | from __future__ import annotations |
| 4 | |
| 5 | import urllib.error |
| 6 | import urllib.request |
| 7 | from dataclasses import dataclass |
| 8 | |
| 9 | _DEFAULT_MAX_BYTES = 100 * 1024 * 1024 |
| 10 | _DEFAULT_TIMEOUT_S = 60.0 |
| 11 | |
| 12 | |
| 13 | @dataclass(frozen=True) |
| 14 | class HttpDownloadResult: |
| 15 | data: bytes |
| 16 | content_type: str |
| 17 | content_length: int | None |
| 18 | |
| 19 | |
| 20 | class HttpDownloadError(RuntimeError): |
| 21 | """Raised when an HTTP download fails or exceeds policy limits.""" |
| 22 | |
| 23 | |
| 24 | class _NoRedirectHandler(urllib.request.HTTPRedirectHandler): |
| 25 | def redirect_request(self, req, fp, code, msg, headers, newurl): |
| 26 | raise HttpDownloadError("下载失败：禁止重定向") |
| 27 | |
| 28 | |
| 29 | def download_http_bytes( |
| 30 | url: str, |
| 31 | *, |
| 32 | max_bytes: int = _DEFAULT_MAX_BYTES, |
| 33 | timeout_s: float = _DEFAULT_TIMEOUT_S, |
| 34 | ) -> HttpDownloadResult: |
| 35 | """Download *url* into memory with SSRF follow-up protections. |
| 36 | |
| 37 | - Does not follow redirects (302 → internal IP bypass). |
| 38 | - Accepts HTTP 200 only. |
| 39 | - Caps response body at *max_bytes* (+1 byte probe for oversize detection). |
| 40 | """ |
| 41 | if max_bytes <= 0: |
| 42 | raise HttpDownloadError("无效的大小限制") |
| 43 | |
| 44 | opener = urllib.request.build_opener(_NoRedirectHandler) |
| 45 | request = urllib.request.Request(url, method="GET") |
| 46 | request.add_header("Accept", "*/*") |
| 47 | |
| 48 | try: |
| 49 | with opener.open(request, timeout=timeout_s) as response: |
| 50 | status = getattr(response, "status", None) or response.getcode() |
| 51 | if status != 200: |
| 52 | raise HttpDownloadError(f"下载失败，HTTP状态码: {status}") |
| 53 | |
| 54 | content_type = response.headers.get("Content-Type") or "application/octet-stream" |
| 55 | raw_length = response.headers.get("Content-Length") |
| 56 | content_length: int | None = None |
| 57 | if raw_length: |
| 58 | try: |
| 59 | content_length = int(raw_length) |
| 60 | except ValueError: |
| 61 | content_length = None |
| 62 | |
| 63 | limited = response.read(max_bytes + 1) |
| 64 | except HttpDownloadError: |
| 65 | raise |
| 66 | except urllib.error.HTTPError as exc: |
| 67 | raise HttpDownloadError(f"下载失败，HTTP状态码: {exc.code}") from exc |
| 68 | except Exception as exc: |
| 69 | raise HttpDownloadError("下载文件失败") from exc |
| 70 | |
| 71 | if len(limited) > max_bytes: |
| 72 | raise HttpDownloadError(f"文件大小超过限制({max_bytes // (1024 * 1024)}MB)") |
| 73 | |
| 74 | return HttpDownloadResult( |
| 75 | data=limited, |
| 76 | content_type=content_type.split(";", 1)[0].strip() or "application/octet-stream", |
| 77 | content_length=content_length, |
| 78 | ) |
| 79 |