"""URL 安全模块单元测试 - 技术债务第二轮：魔数校验 + 重定向每跳校验.""" from __future__ import annotations import os from pathlib import Path import pytest # ══════════════════════════════════════════════════════════════════════════════ # 魔数校验测试 # ══════════════════════════════════════════════════════════════════════════════ class TestMagicNumberValidation: """文件头魔数校验测试.""" def test_png_magic_passes(self, tmp_path: Path): """PNG 魔数正确应通过校验.""" from shared.url_security import _validate_magic_number f = tmp_path / "test.png" f.write_bytes(b"\x89PNG\r\n\x1a\n" + b"\x00" * 100) _validate_magic_number(str(f), {"image/png"}) # 不抛异常即通过 def test_jpeg_magic_passes(self, tmp_path: Path): """JPEG 魔数正确应通过校验.""" from shared.url_security import _validate_magic_number f = tmp_path / "test.jpg" f.write_bytes(b"\xff\xd8\xff\xe0" + b"\x00" * 100) _validate_magic_number(str(f), {"image/jpeg", "image/png"}) def test_gif_magic_passes(self, tmp_path: Path): """GIF 魔数正确应通过校验.""" from shared.url_security import _validate_magic_number f = tmp_path / "test.gif" f.write_bytes(b"GIF89a" + b"\x00" * 100) _validate_magic_number(str(f), {"image/gif"}) def test_mp3_magic_id3_passes(self, tmp_path: Path): """MP3 ID3v2 标签魔数应通过校验.""" from shared.url_security import _validate_magic_number f = tmp_path / "test.mp3" f.write_bytes(b"ID3\x03\x00\x00\x00\x00\x00\x00" + b"\x00" * 100) _validate_magic_number(str(f), {"audio/mpeg"}) def test_mp3_magic_frame_passes(self, tmp_path: Path): """MP3 frame sync 魔数应通过校验.""" from shared.url_security import _validate_magic_number f = tmp_path / "test.mp3" f.write_bytes(b"\xff\xfb\x90\x00" + b"\x00" * 100) _validate_magic_number(str(f), {"audio/mpeg"}) def test_wav_magic_passes(self, tmp_path: Path): """WAV RIFF+WAVE 魔数应通过校验.""" from shared.url_security import _validate_magic_number f = tmp_path / "test.wav" header = b"RIFF" + b"\x24\x00\x00\x00" + b"WAVE" + b"fmt " + b"\x00" * 100 f.write_bytes(header) _validate_magic_number(str(f), {"audio/wav"}) def test_mp4_magic_passes(self, tmp_path: Path): """MP4 ftyp 魔数应通过校验.""" from shared.url_security import _validate_magic_number f = tmp_path / "test.mp4" # ftyp box: size(4) + 'ftyp'(4) + major_brand(4) + ... f.write_bytes(b"\x00\x00\x00\x20ftypisom" + b"\x00" * 100) _validate_magic_number(str(f), {"video/mp4"}) def test_webp_magic_passes(self, tmp_path: Path): """WebP RIFF+WEBP 魔数应通过校验.""" from shared.url_security import _validate_magic_number f = tmp_path / "test.webp" f.write_bytes(b"RIFF\x00\x00\x00\x00WEBP" + b"\x00" * 100) _validate_magic_number(str(f), {"image/webp"}) def test_wrong_magic_raises(self, tmp_path: Path): """魔数不匹配应抛出 UrlSecurityError.""" from shared.url_security import UrlSecurityError, _validate_magic_number f = tmp_path / "fake.png" f.write_bytes(b"NOT_A_PNG_FILE!!!" + b"\x00" * 100) with pytest.raises(UrlSecurityError, match="魔数"): _validate_magic_number(str(f), {"image/png", "image/jpeg"}) def test_text_as_png_raises(self, tmp_path: Path): """纯文本伪装成 PNG 应被拦截.""" from shared.url_security import UrlSecurityError, _validate_magic_number f = tmp_path / "fake.png" f.write_text("not an image", encoding="utf-8") with pytest.raises(UrlSecurityError): _validate_magic_number(str(f), {"image/png"}) def test_empty_file_raises(self, tmp_path: Path): """空文件应抛出异常.""" from shared.url_security import UrlSecurityError, _validate_magic_number f = tmp_path / "empty.png" f.write_bytes(b"") with pytest.raises(UrlSecurityError, match="为空"): _validate_magic_number(str(f), {"image/png"}) def test_unknown_mime_skipped(self, tmp_path: Path): """未知 MIME 类型没有对应魔数，应跳过校验不阻断.""" from shared.url_security import _validate_magic_number f = tmp_path / "test.xyz" f.write_bytes(b"random garbage data here") # 没有已知魔数的 MIME，跳过校验 _validate_magic_number(str(f), {"application/x-custom-format"}) def test_multiple_allowed_types_one_matches(self, tmp_path: Path): """多个允许类型，只要有一个匹配就通过.""" from shared.url_security import _validate_magic_number f = tmp_path / "test" f.write_bytes(b"\x89PNG\r\n\x1a\n" + b"\x00" * 50) _validate_magic_number(str(f), {"image/jpeg", "image/png", "image/gif"}) def test_multiple_allowed_types_none_match(self, tmp_path: Path): """多个允许类型都不匹配应抛异常.""" from shared.url_security import UrlSecurityError, _validate_magic_number f = tmp_path / "test" f.write_bytes(b"RIFF\x00\x00\x00\x00WAVE" + b"\x00" * 50) with pytest.raises(UrlSecurityError): _validate_magic_number(str(f), {"image/png", "image/jpeg", "image/gif"}) def test_flac_magic_passes(self, tmp_path: Path): """FLAC 魔数应通过校验.""" from shared.url_security import _validate_magic_number f = tmp_path / "test.flac" f.write_bytes(b"fLaC" + b"\x00" * 100) _validate_magic_number(str(f), {"audio/flac"}) def test_ogg_magic_passes(self, tmp_path: Path): """OGG 魔数应通过校验.""" from shared.url_security import _validate_magic_number f = tmp_path / "test.ogg" f.write_bytes(b"OggS\x00\x02\x00\x00" + b"\x00" * 100) _validate_magic_number(str(f), {"audio/ogg"}) def test_bmp_magic_passes(self, tmp_path: Path): """BMP 魔数应通过校验.""" from shared.url_security import _validate_magic_number f = tmp_path / "test.bmp" f.write_bytes(b"BM\x00\x00\x00\x00" + b"\x00" * 100) _validate_magic_number(str(f), {"image/bmp"}) # ══════════════════════════════════════════════════════════════════════════════ # safe_download_file 魔数校验集成测试 # ══════════════════════════════════════════════════════════════════════════════ class TestSafeDownloadMagicIntegration: """safe_download_file 集成魔数校验测试.""" def test_download_with_mime_and_magic_match(self, tmp_path: Path): """MIME 匹配 + 魔数匹配，下载成功.""" from unittest.mock import MagicMock, patch from shared.url_security import safe_download_file png_data = b"\x89PNG\r\n\x1a\n" + b"\x00" * 200 class FakeResp: headers = {"Content-Type": "image/png", "Content-Length": str(len(png_data))} def read(self, n): if not hasattr(self, "_pos"): self._pos = 0 chunk = png_data[self._pos : self._pos + n] self._pos += len(chunk) return chunk def close(self): pass class FakeOpener: def open(self, req, timeout=None): return FakeResp() with ( patch("shared.url_security.validate_url_safety", side_effect=lambda u, **kw: u), patch("shared.url_security.urllib.request.build_opener", return_value=FakeOpener()), ): dest = str(tmp_path / "out.png") size = safe_download_file( "https://example.com/test.png", dest, allowed_mime_types={"image/png"}, purpose="test", ) assert size == len(png_data) with open(dest, "rb") as f: assert f.read() == png_data def test_download_mime_match_but_magic_mismatch_raises(self, tmp_path: Path): """Content-Type 声明是 PNG 但实际文件是 HTML，应被魔数校验拦截.""" from unittest.mock import patch from shared.url_security import UrlSecurityError, safe_download_file fake_data = b"not really a png" class FakeResp: headers = {"Content-Type": "image/png", "Content-Length": str(len(fake_data))} def read(self, n): if not hasattr(self, "_pos"): self._pos = 0 chunk = fake_data[self._pos : self._pos + n] self._pos += len(chunk) return chunk def close(self): pass class FakeOpener: def open(self, req, timeout=None): return FakeResp() with ( patch("shared.url_security.validate_url_safety", side_effect=lambda u, **kw: u), patch("shared.url_security.urllib.request.build_opener", return_value=FakeOpener()), ): dest = str(tmp_path / "out.png") with pytest.raises(UrlSecurityError, match="魔数"): safe_download_file( "https://example.com/fake.png", dest, allowed_mime_types={"image/png"}, purpose="test", ) def test_download_no_mime_check_skips_magic(self, tmp_path: Path): """不传 allowed_mime_types 时不做 MIME 校验也不做魔数校验.""" from unittest.mock import patch from shared.url_security import safe_download_file data = b"any random content here" class FakeResp: headers = {"Content-Type": "application/octet-stream"} def read(self, n): if not hasattr(self, "_pos"): self._pos = 0 chunk = data[self._pos : self._pos + n] self._pos += len(chunk) return chunk def close(self): pass class FakeOpener: def open(self, req, timeout=None): return FakeResp() with ( patch("shared.url_security.validate_url_safety", side_effect=lambda u, **kw: u), patch("shared.url_security.urllib.request.build_opener", return_value=FakeOpener()), ): dest = str(tmp_path / "out.bin") size = safe_download_file( "https://example.com/file.bin", dest, purpose="test", ) assert size == len(data) # ══════════════════════════════════════════════════════════════════════════════ # _verify_url_accessible 重定向每跳校验测试 # ══════════════════════════════════════════════════════════════════════════════ class TestVerifyUrlRedirectValidation: """URL 可访问性校验 - 重定向每跳 SSRF 校验测试. 直接复制核心逻辑进行单元测试，避免导入 generation 模块触发 DB 连接。 逻辑与 generation.py 中的 _verify_url_accessible 完全一致。 """ @staticmethod def _verify_url_accessible(url, timeout=10.0, retries=0, max_redirects=5): """从 generation.py 复制的核心逻辑，用于单元测试.""" import time import urllib.request from urllib.parse import urljoin from shared.url_security import UrlSecurityError, validate_url_safety class NoRedirect(urllib.request.HTTPRedirectHandler): def redirect_request(self, req, fp, code, msg, headers, newurl): # noqa: N802 return None last_error = None def _do_verify(current_url): redirect_count = 0 url_being_checked = current_url opener = urllib.request.build_opener(NoRedirect()) while redirect_count <= max_redirects: safe_url = validate_url_safety(url_being_checked, purpose="url_verify") req = urllib.request.Request(safe_url, method="HEAD") req.add_header("User-Agent", "xiaoxia-saas-worker/1.0") with opener.open(req, timeout=timeout): # 简化：进入 with 块即表示 2xx（3xx 被 NoRedirect 拦截为 HTTPError） return True raise Exception("unreachable") import urllib.error for attempt in range(1 + retries): try: # 用 try/except 手动处理重定向 redirect_count = 0 current = url opener = urllib.request.build_opener(NoRedirect()) while redirect_count <= max_redirects: safe_url = validate_url_safety(current, purpose="url_verify") req = urllib.request.Request(safe_url, method="HEAD") req.add_header("User-Agent", "xiaoxia-saas-worker/1.0") try: with opener.open(req, timeout=timeout) as resp: if 200 <= resp.status < 300: return True if resp.status < 400: return True last_error = Exception(f"HTTP {resp.status}") except urllib.error.HTTPError as e: if 300 <= e.code < 400 and e.headers.get("Location"): if redirect_count >= max_redirects: raise Exception(f"重定向次数超过上限 ({max_redirects})") from e location = e.headers["Location"] current = urljoin(safe_url, location) redirect_count += 1 continue last_error = Exception(f"HTTP {e.code}") break else: raise Exception(f"重定向次数超过上限 ({max_redirects})") except Exception as e: last_error = e if attempt < retries: time.sleep(0) return False def test_simple_200_ok(self): """普通 200 响应应返回 True.""" import urllib.error from unittest.mock import patch class FakeResp: status = 200 headers = {} def __enter__(self): return self def __exit__(self, *args): pass class FakeOpener: def open(self, req, timeout=None): return FakeResp() with ( patch("urllib.request.build_opener", return_value=FakeOpener()), patch("shared.url_security.validate_url_safety", side_effect=lambda u, **kw: u), ): result = self._verify_url_accessible("https://example.com/file.mp4", retries=0) assert result is True def test_redirect_to_internal_ip_blocked(self): """重定向到内网 IP 应被拦截（返回 False）.""" import urllib.error from unittest.mock import patch from shared.url_security import UrlSecurityError call_count = 0 class FakeHTTPError(urllib.error.HTTPError): def __init__(self): pass # 用 validate_url_safety 来模拟拦截 def fake_validate(url, **kwargs): if "127.0.0.1" in url: raise UrlSecurityError("内网IP禁止访问") return url class FakeOpener: def open(self, req, timeout=None): nonlocal call_count call_count += 1 # 第一次请求返回 302 raise urllib.error.HTTPError( req.full_url, 302, "Found", {"Location": "http://127.0.0.1/internal"}, None ) with ( patch("urllib.request.build_opener", return_value=FakeOpener()), patch("shared.url_security.validate_url_safety", side_effect=fake_validate), ): result = self._verify_url_accessible("https://example.com/redirect", retries=0) assert result is False assert call_count == 1 # 只请求了第一次，第二次跳转在校验阶段就被拦了 def test_redirect_count_exceeded(self): """超过最大重定向次数应返回 False.""" import urllib.error from unittest.mock import patch call_count = 0 class FakeOpener: def open(self, req, timeout=None): nonlocal call_count call_count += 1 raise urllib.error.HTTPError(req.full_url, 302, "Found", {"Location": "https://example.com/next"}, None) with ( patch("urllib.request.build_opener", return_value=FakeOpener()), patch("shared.url_security.validate_url_safety", side_effect=lambda u, **kw: u), ): result = self._verify_url_accessible( "https://example.com/start", retries=0, max_redirects=3, ) assert result is False assert call_count == 4 # 初始 + 3次跳转 = 4次请求 def test_redirect_chain_valid(self): """合法的重定向链（都是公网域名）应返回 True.""" import urllib.error from unittest.mock import patch step = 0 class FakeResp: status = 200 headers = {} def __enter__(self): return self def __exit__(self, *args): pass class FakeOpener: def open(self, req, timeout=None): nonlocal step step += 1 if step == 1: raise urllib.error.HTTPError( req.full_url, 302, "Found", {"Location": "https://cdn.example.com/final.mp4"}, None ) return FakeResp() with ( patch("urllib.request.build_opener", return_value=FakeOpener()), patch("shared.url_security.validate_url_safety", side_effect=lambda u, **kw: u), ): result = self._verify_url_accessible( "https://example.com/redirect", retries=0, max_redirects=5, ) assert result is True assert step == 2 def test_404_returns_false(self): """404 应返回 False.""" import urllib.error from unittest.mock import patch class FakeOpener: def open(self, req, timeout=None): raise urllib.error.HTTPError(req.full_url, 404, "Not Found", {}, None) with ( patch("urllib.request.build_opener", return_value=FakeOpener()), patch("shared.url_security.validate_url_safety", side_effect=lambda u, **kw: u), ): result = self._verify_url_accessible("https://example.com/nonexistent", retries=0) assert result is False