xiaoxia
531aacb57e
CI/CD Pipeline / Validate Code Quality And Tests (push) Has been cancelled
CI/CD Pipeline / Unit Tests (push) Has been cancelled
CI/CD Pipeline / Integration Tests (push) Has been cancelled
CI/CD Pipeline / Frontend Lint (push) Has been cancelled
CI/CD Pipeline / Build Staging API Image (push) Has been cancelled
CI/CD Pipeline / Build Staging Worker Image (push) Has been cancelled
CI/CD Pipeline / Build Staging Web Image (push) Has been cancelled
CI/CD Pipeline / Deploy Staging (Watchtower auto-deploy) (push) Has been cancelled
CI/CD Pipeline / Staging E2E Tests (push) Has been cancelled
CI/CD Pipeline / Staging API Integration Tests (push) Has been cancelled
CI/CD Pipeline / Build Production API Image (push) Has been cancelled
CI/CD Pipeline / Build Production Worker Image (push) Has been cancelled
CI/CD Pipeline / Build Production Web Image (push) Has been cancelled
CI/CD Pipeline / Deploy Production (push) Has been cancelled
CI/CD Pipeline / Production Browser E2E (push) Has been cancelled
528 lines
20 KiB
Python
528 lines
20 KiB
Python
"""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("<html>not an image</html>", 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"<html>not really a png</html>"
|
||
|
||
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
|