"""通用 HTTP Range 代理。 供浏览器、WebDAV、STRM、文件下载这些需要本地代理的场景共用。 """ import asyncio import mimetypes from dataclasses import replace from datetime import datetime from email.utils import format_datetime, parsedate_to_datetime from typing import Optional, Tuple from urllib.parse import quote import aiohttp from fastapi import Request from fastapi.responses import Response, StreamingResponse from core.base import FileItem from core.driver_service import ( build_upstream_download_headers, get_driver_proxy_part_size, resolve_download, ) _proxy_session: Optional[aiohttp.ClientSession] = None def get_range_proxy_session() -> aiohttp.ClientSession: global _proxy_session if _proxy_session is None or _proxy_session.closed: timeout = aiohttp.ClientTimeout(total=None, connect=30, sock_read=300) connector = aiohttp.TCPConnector(limit=1000, limit_per_host=100, keepalive_timeout=60) cookie_jar = aiohttp.DummyCookieJar() _proxy_session = aiohttp.ClientSession(timeout=timeout, connector=connector, cookie_jar=cookie_jar) return _proxy_session async def close_range_proxy_session() -> None: global _proxy_session session = _proxy_session _proxy_session = None if session and not session.closed: await session.close() # ---- MIME / Content-Type ---- _FALLBACK_MIME_TYPES = { "mp4": "video/mp4", "m4v": "video/x-m4v", "iso": "application/x-iso9660-image", "mkv": "video/x-matroska", "ts": "video/mp2t", "m2ts": "video/mp2t", "mts": "video/mp2t", "flv": "video/x-flv", "webm": "video/webm", "mov": "video/quicktime", "avi": "video/x-msvideo", "wmv": "video/x-ms-wmv", "rmvb": "video/vnd.rn-realmedia-vbr", "rm": "application/vnd.rn-realmedia", "3gp": "video/3gpp", "3g2": "video/3gpp2", "asf": "video/x-ms-asf", "vob": "video/dvd", "f4v": "video/mp4", "mpd": "application/dash+xml", "m3u8": "application/vnd.apple.mpegurl", "mp3": "audio/mpeg", "aac": "audio/aac", "flac": "audio/flac", "wav": "audio/wav", "ogg": "audio/ogg", "oga": "audio/ogg", "opus": "audio/ogg", "m4a": "audio/mp4", "wma": "audio/x-ms-wma", "ape": "audio/x-ape", "alac": "audio/alac", } def guess_content_type(file_name: str) -> str: name = file_name or "" mime_type, _ = mimetypes.guess_type(name) if mime_type: return mime_type if "." in name: ext = name.rsplit(".", 1)[-1].strip().lower() if ext in _FALLBACK_MIME_TYPES: return _FALLBACK_MIME_TYPES[ext] return "application/octet-stream" def is_stream_preview_type(file_name: str) -> bool: content_type = guess_content_type(file_name) return content_type.startswith("video/") or content_type.startswith("audio/") def build_stable_etag(file_info: FileItem) -> str: modified_part = "" if file_info.modified: try: modified_part = f"-{int(file_info.modified.timestamp())}" except Exception: modified_part = "" safe_id = quote(str(file_info.id), safe="") return f"\"{safe_id}-{file_info.size}{modified_part}\"" def build_proxy_file_info( file_id: str, *, file_name: str, file_size: int, template: Optional[FileItem] = None, ) -> FileItem: modified = getattr(template, "modified", None) if template else None created = getattr(template, "created", None) if template else None name = file_name or (getattr(template, "name", None) if template else None) or f"file_{file_id}" return FileItem( id=file_id, name=name, path=getattr(template, "path", "") if template else "", size=int(file_size or 0), is_dir=False, modified=modified, created=created, ) def build_proxy_file_info_from_download(file_id: str, download) -> Tuple[Optional[FileItem], str, int]: raw_file_info = getattr(download, "file_info", None) file_size = int(getattr(download, "file_size", 0) or 0) raw_size = int(getattr(raw_file_info, "size", 0) or 0) if raw_file_info else 0 if raw_file_info and not raw_file_info.is_dir and raw_size > 0: return raw_file_info, download.download_url, raw_size effective_size = file_size if file_size > 0 else raw_size if effective_size <= 0: return raw_file_info, download.download_url, 0 file_info = build_proxy_file_info( file_id, file_name=getattr(download, "file_name", "") or f"file_{file_id}", file_size=effective_size, template=raw_file_info, ) return file_info, download.download_url, effective_size # ---- Range 解析 / 条件请求 ---- def parse_single_range(range_header: str, total_size: int) -> Tuple[int, int]: if total_size <= 0: raise ValueError("empty") header_value = (range_header or "").strip() if not header_value.startswith("bytes="): raise ValueError("unit") ranges_str = header_value[len("bytes="):].strip() parts = [part.strip() for part in ranges_str.split(",") if part.strip()] if len(parts) != 1: raise ValueError("multiple") part = parts[0] if "-" not in part: raise ValueError("format") start_str, end_str = part.split("-", 1) if start_str == "": suffix_length = int(end_str) if end_str else 0 if suffix_length <= 0: raise ValueError("suffix") if suffix_length >= total_size: return 0, total_size - 1 return total_size - suffix_length, total_size - 1 start = int(start_str) if start < 0 or start >= total_size: raise ValueError("start") if end_str == "": return start, total_size - 1 end = int(end_str) if end < start: raise ValueError("order") if end >= total_size: end = total_size - 1 return start, end def evaluate_if_range(request: Request, etag: str, last_modified: Optional[datetime]) -> bool: value = (request.headers.get("If-Range") or "").strip() if not value: return True if value.startswith('"'): return value == etag try: if_range_dt = parsedate_to_datetime(value) if if_range_dt.tzinfo: if_range_dt = if_range_dt.astimezone(tz=None).replace(tzinfo=None) if last_modified and last_modified.replace(microsecond=0) <= if_range_dt.replace(microsecond=0): return True return False except Exception: return False def match_if_none_match(request: Request, etag: str) -> bool: value = (request.headers.get("If-None-Match") or "").strip() if not value: return False if value == "*": return True candidates = [item.strip() for item in value.split(",") if item.strip()] return etag in candidates def match_if_modified_since(request: Request, last_modified: Optional[datetime]) -> bool: value = (request.headers.get("If-Modified-Since") or "").strip() if not value or not last_modified: return False try: since = parsedate_to_datetime(value) if since.tzinfo: since = since.astimezone(tz=None).replace(tzinfo=None) return last_modified.replace(microsecond=0) <= since.replace(microsecond=0) except Exception: return False def format_http_datetime(value: Optional[datetime]) -> Optional[str]: if not value: return None try: return format_datetime(value, usegmt=True) except Exception: return None DEFAULT_PROXY_PART_SIZE = 8 * 1024 * 1024 def max_slice_size(file_size: Optional[int] = None) -> int: return DEFAULT_PROXY_PART_SIZE def resolve_chunk_cap(driver) -> int: declared = get_driver_proxy_part_size(driver) return declared if declared > 0 else DEFAULT_PROXY_PART_SIZE def stream_chunk_size(file_size: Optional[int]) -> int: if not file_size: return 262144 if file_size > 42949672960: return 2097152 if file_size > 10737418240: return 1048576 if file_size > 2147483648: return 524288 return 262144 _CLIENT_FORWARD_HEADERS = ( "origin", "accept-language", "sec-fetch-site", "sec-fetch-mode", "sec-fetch-dest", "sec-ch-ua", "sec-ch-ua-mobile", "sec-ch-ua-platform", ) async def open_upstream_get( *, driver, file_id: str, user_agent: str, initial_url: str, file_info: Optional[FileItem], range_header: Optional[str], client_headers: Optional[dict] = None, upstream_headers_override: Optional[dict] = None, ) -> aiohttp.ClientResponse: session = get_range_proxy_session() current_url = initial_url current_headers_override = upstream_headers_override resp: Optional[aiohttp.ClientResponse] = None for attempt in range(2): headers = await build_upstream_download_headers( driver, file_id, user_agent, range_header=range_header, prefer_identity=True, headers_override=current_headers_override, ) if client_headers: for key in _CLIENT_FORWARD_HEADERS: value = client_headers.get(key, "") if value: headers.setdefault(key, value) headers.pop("Cache-Control", None) resp = await session.get(current_url, headers=headers) if resp.status in {401, 403} and attempt == 0: resp.release() refreshed = await resolve_download(driver, file_id, user_agent, file_info=file_info) if refreshed.download_url: current_url = refreshed.download_url current_headers_override = refreshed.headers continue return resp assert resp is not None return resp def _base_headers( file_info: FileItem, content_type: str, etag: str, content_disposition: str = "", ) -> dict: headers = { "Content-Type": content_type, "Accept-Ranges": "bytes", "ETag": etag, } last_modified = format_http_datetime(file_info.modified) if last_modified: headers["Last-Modified"] = last_modified if content_disposition: headers["Content-Disposition"] = content_disposition return headers async def _stream_chunked_range( *, driver, file_id: str, user_agent: str, initial_url: str, file_info: Optional[FileItem], range_start: int, range_end: int, chunk_cap: int, client_headers: Optional[dict] = None, upstream_headers_override: Optional[dict] = None, ): cur = range_start while cur <= range_end: slice_end = min(cur + chunk_cap - 1, range_end) upstream_range = f"bytes={cur}-{slice_end}" resp = await open_upstream_get( driver=driver, file_id=file_id, user_agent=user_agent, initial_url=initial_url, file_info=file_info, range_header=upstream_range, client_headers=client_headers, upstream_headers_override=upstream_headers_override, ) if resp.status >= 400: try: resp.release() except Exception: pass return try: async for chunk in resp.content.iter_chunked(1048576): yield chunk except (asyncio.CancelledError, ConnectionResetError, BrokenPipeError): try: resp.release() except Exception: pass return finally: try: resp.release() except Exception: pass cur = slice_end + 1 # ---- HEAD ---- def _parse_total_from_content_range(value: str) -> int: v = (value or "").strip() if "/" not in v: return 0 tail = v.split("/")[-1].strip() if tail == "*": return 0 try: return int(tail) except ValueError: return 0 async def _probe_file_size_via_range0( *, driver, file_id: str, user_agent: str, initial_url: str, file_info: FileItem, client_headers: Optional[dict] = None, upstream_headers_override: Optional[dict] = None, ) -> int: resp = await open_upstream_get( driver=driver, file_id=file_id, user_agent=user_agent, initial_url=initial_url, file_info=file_info, range_header="bytes=0-0", client_headers=client_headers, upstream_headers_override=upstream_headers_override, ) try: if resp.status == 206: return _parse_total_from_content_range(resp.headers.get("Content-Range", "")) finally: try: resp.release() except Exception: pass return 0 async def _serve_forwarded_upstream_range( *, driver, file_id: str, file_info: FileItem, etag: str, user_agent: str, range_header: str, initial_url: str, content_type: str, content_disposition: str, client_headers: Optional[dict] = None, upstream_headers_override: Optional[dict] = None, ) -> Response: """总大小未知时把客户端 Range 交给 CDN；返回 206 则透传（ISO 等依赖随机 Range）。""" rh = (range_header or "").strip() if not rh.lower().startswith("bytes="): return Response(status_code=400, content="Bad Range unit") ranges_blob = rh[6:].strip() if "," in ranges_blob: return Response(status_code=416, content="Multiple ranges not supported") upstream = await open_upstream_get( driver=driver, file_id=file_id, user_agent=user_agent, initial_url=initial_url, file_info=file_info, range_header=rh, client_headers=client_headers, upstream_headers_override=upstream_headers_override, ) st = upstream.status cr = (upstream.headers.get("Content-Range") or "").strip() cl = (upstream.headers.get("Content-Length") or "").strip() chunk_size = stream_chunk_size(0) def base_hdrs() -> dict: h = _base_headers(file_info, content_type, etag, content_disposition) if cr: h["Content-Range"] = cr if cl: h["Content-Length"] = cl return h if st == 206: async def gen206(): try: async for chunk in upstream.content.iter_chunked(chunk_size): yield chunk finally: try: upstream.release() except Exception: pass return StreamingResponse( gen206(), status_code=206, headers=base_hdrs(), media_type=content_type ) # 少数 CDN 对 Range 仍回 200，但带 Content-Range（视作分片成功） if st == 200 and cr and _parse_total_from_content_range(cr) > 0: async def gen200partial(): try: async for chunk in upstream.content.iter_chunked(chunk_size): yield chunk finally: try: upstream.release() except Exception: pass hdrs = base_hdrs() hdrs["Content-Range"] = cr if cl: hdrs["Content-Length"] = cl return StreamingResponse( gen200partial(), status_code=206, headers=hdrs, media_type=content_type ) try: upstream.release() except Exception: pass return Response( status_code=502, content="Upstream did not return partial content for range request", ) def build_head_response( file_info: FileItem, *, content_type: Optional[str] = None, content_disposition: str = "", ) -> Response: """HEAD 响应：只回元数据 + Accept-Ranges，让 NDM/IDM 探测到「能续传」。""" total_size = int(file_info.size or 0) ct = content_type or guess_content_type(file_info.name or "") etag = build_stable_etag(file_info) headers = _base_headers(file_info, ct, etag, content_disposition) if total_size > 0: headers["Content-Length"] = str(total_size) return Response(status_code=200, headers=headers, media_type=ct) # ---- 主入口：serve_range_proxy ---- async def serve_range_proxy( *, driver, file_id: str, file_info: FileItem, request: Request, initial_url: str, content_disposition: str = "", user_agent_override: Optional[str] = None, upstream_headers_override: Optional[dict] = None, ) -> Response: if file_info.is_dir: return Response(status_code=400, content="Cannot stream a directory") total_size = int(file_info.size or 0) etag = build_stable_etag(file_info) content_type = guess_content_type(file_info.name or "") user_agent = user_agent_override if user_agent_override is not None else request.headers.get("User-Agent", "") client_hdrs = dict(request.headers) if request.method == "HEAD": effective_info = file_info if int(file_info.size or 0) <= 0: probed = await _probe_file_size_via_range0( driver=driver, file_id=file_id, user_agent=user_agent, initial_url=initial_url, file_info=file_info, client_headers=client_hdrs, upstream_headers_override=upstream_headers_override, ) if probed > 0: effective_info = replace(file_info, size=probed) return build_head_response( effective_info, content_type=content_type, content_disposition=content_disposition ) raw_range_header = request.headers.get("Range") # 304 / If-None-Match / If-Modified-Since if not raw_range_header and ( match_if_none_match(request, etag) or match_if_modified_since(request, file_info.modified) ): headers = {"ETag": etag} last_modified = format_http_datetime(file_info.modified) if last_modified: headers["Last-Modified"] = last_modified if content_disposition: headers["Content-Disposition"] = content_disposition return Response(status_code=304, headers=headers) # If-Range 不匹配 → 当作无 Range 处理 if raw_range_header and not evaluate_if_range(request, etag, file_info.modified): raw_range_header = None if raw_range_header and total_size > 0: return await _serve_range( driver=driver, file_id=file_id, file_info=file_info, etag=etag, user_agent=user_agent, range_header=raw_range_header, total_size=total_size, initial_url=initial_url, content_type=content_type, content_disposition=content_disposition, client_headers=client_hdrs, upstream_headers_override=upstream_headers_override, ) if raw_range_header and total_size <= 0: return await _serve_forwarded_upstream_range( driver=driver, file_id=file_id, file_info=file_info, etag=etag, user_agent=user_agent, range_header=raw_range_header, initial_url=initial_url, content_type=content_type, content_disposition=content_disposition, client_headers=client_hdrs, upstream_headers_override=upstream_headers_override, ) return await _serve_full( driver=driver, file_id=file_id, file_info=file_info, etag=etag, user_agent=user_agent, total_size=total_size, initial_url=initial_url, content_type=content_type, content_disposition=content_disposition, client_headers=client_hdrs, upstream_headers_override=upstream_headers_override, ) async def _serve_range( *, driver, file_id: str, file_info: FileItem, etag: str, user_agent: str, range_header: str, total_size: int, initial_url: str, content_type: str, content_disposition: str, client_headers: Optional[dict] = None, upstream_headers_override: Optional[dict] = None, ) -> Response: try: start, end = parse_single_range(range_header, total_size) except Exception: headers = {"Content-Range": f"bytes */{total_size}"} if total_size > 0 else {} return Response(status_code=416, content="Invalid range request", headers=headers) chunk_cap = resolve_chunk_cap(driver) headers = _base_headers(file_info, content_type, etag, content_disposition) headers["Content-Range"] = f"bytes {start}-{end}/{total_size}" headers["Content-Length"] = str(end - start + 1) return StreamingResponse( _stream_chunked_range( driver=driver, file_id=file_id, user_agent=user_agent, initial_url=initial_url, file_info=file_info, range_start=start, range_end=end, chunk_cap=chunk_cap, client_headers=client_headers, upstream_headers_override=upstream_headers_override, ), status_code=206, headers=headers, media_type=content_type, ) async def _serve_full( *, driver, file_id: str, file_info: FileItem, etag: str, user_agent: str, total_size: int, initial_url: str, content_type: str, content_disposition: str, client_headers: Optional[dict] = None, upstream_headers_override: Optional[dict] = None, ) -> Response: headers = _base_headers(file_info, content_type, etag, content_disposition) if total_size <= 0: upstream = await open_upstream_get( driver=driver, file_id=file_id, user_agent=user_agent, initial_url=initial_url, file_info=file_info, range_header=None, client_headers=client_headers, upstream_headers_override=upstream_headers_override, ) if upstream.status >= 400: status = upstream.status upstream.release() return Response(status_code=status, content=f"Upstream returned {status}") chunk_size = stream_chunk_size(total_size) async def stream_unknown_size(): try: async for chunk in upstream.content.iter_chunked(chunk_size): yield chunk except (asyncio.CancelledError, ConnectionResetError, BrokenPipeError): return finally: try: upstream.release() except Exception: pass upstream_length = (upstream.headers.get("Content-Length") or "").strip() if upstream_length.isdigit(): headers["Content-Length"] = upstream_length return StreamingResponse(stream_unknown_size(), status_code=200, headers=headers, media_type=content_type) # 已知总大小：分片拉，对外 Content-Length=total chunk_cap = resolve_chunk_cap(driver) headers["Content-Length"] = str(total_size) return StreamingResponse( _stream_chunked_range( driver=driver, file_id=file_id, user_agent=user_agent, initial_url=initial_url, file_info=file_info, range_start=0, range_end=total_size - 1, chunk_cap=chunk_cap, client_headers=client_headers, upstream_headers_override=upstream_headers_override, ), status_code=200, headers=headers, media_type=content_type, ) # ---- 兼容老调用点（api/strm.py 旧版本曾 import） ---- async def assemble_chunked_proxy( *, driver, file_id: str, user_agent: str, download_url: str, download_size: int, client_headers: Optional[dict] = None, chunk_cap: Optional[int] = None, ): """保留旧接口：等价于 _stream_chunked_range(0, download_size-1)。""" if download_size <= 0: async def _empty(): if False: yield b"" return _empty() cap = chunk_cap or max_slice_size(download_size) return _stream_chunked_range( driver=driver, file_id=file_id, user_agent=user_agent, initial_url=download_url, file_info=None, range_start=0, range_end=download_size - 1, chunk_cap=cap, client_headers=client_headers, )