test(p3-1): 第十波 - url_security 安全模块单测 72个 #714
+553
-252
@@ -1,296 +1,597 @@
|
||||
"""URL 安全校验工具单元测试 — SSRF 防护."""
|
||||
"""
|
||||
url_security URL安全校验单元测试
|
||||
|
||||
from __future__ import annotations
|
||||
覆盖：
|
||||
- validate_url_safety: scheme/主机/端口/SSRF/内网域名/白名单
|
||||
- is_url_safe: 便捷函数
|
||||
- UrlSecurityError / NoRedirectHandler
|
||||
- _validate_magic_number: 文件魔数校验
|
||||
- safe_download_file / safe_download_bytes: mock 网络测试
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import unittest
|
||||
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", "apps", "worker"))
|
||||
|
||||
import shutil
|
||||
import tempfile
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from video_processing.url_security import ( # noqa: E402
|
||||
import pytest
|
||||
|
||||
from packages.shared.url_security import (
|
||||
ALLOWED_AUDIO_MIME_TYPES,
|
||||
ALLOWED_IMAGE_MIME_TYPES,
|
||||
ALLOWED_PORTS,
|
||||
ALLOWED_SCHEMES,
|
||||
MAX_URL_LENGTH,
|
||||
NoRedirectHandler,
|
||||
UrlSecurityError,
|
||||
_check_internal_hostnames,
|
||||
_is_trusted_domain,
|
||||
_validate_magic_number,
|
||||
is_url_safe,
|
||||
safe_download_bytes,
|
||||
safe_download_file,
|
||||
validate_url_safety,
|
||||
)
|
||||
|
||||
|
||||
class TestUrlSecurityValidation(unittest.TestCase):
|
||||
"""URL 安全校验测试."""
|
||||
|
||||
# ── Scheme 白名单 ──────────────────────────────────────────────────────
|
||||
|
||||
def test_http_scheme_allowed(self):
|
||||
"""HTTP scheme 应该被允许."""
|
||||
result = validate_url_safety("http://example.com/test", purpose="test")
|
||||
self.assertEqual(result, "http://example.com/test")
|
||||
|
||||
def test_https_scheme_allowed(self):
|
||||
"""HTTPS scheme 应该被允许."""
|
||||
result = validate_url_safety("https://example.com/test", purpose="test")
|
||||
self.assertEqual(result, "https://example.com/test")
|
||||
|
||||
def test_file_scheme_rejected(self):
|
||||
"""file:// scheme 应该被拒绝."""
|
||||
with self.assertRaises(UrlSecurityError):
|
||||
validate_url_safety("file:///etc/passwd", purpose="test")
|
||||
|
||||
def test_ftp_scheme_rejected(self):
|
||||
"""ftp:// scheme 应该被拒绝."""
|
||||
with self.assertRaises(UrlSecurityError):
|
||||
validate_url_safety("ftp://example.com/test", purpose="test")
|
||||
|
||||
def test_empty_scheme_rejected(self):
|
||||
"""空 scheme 应该被拒绝."""
|
||||
with self.assertRaises(UrlSecurityError):
|
||||
validate_url_safety("example.com/test", purpose="test")
|
||||
|
||||
# ── 端口白名单 ────────────────────────────────────────────────────────
|
||||
|
||||
def test_port_80_allowed(self):
|
||||
"""端口 80 应该被允许."""
|
||||
# 80端口是默认HTTP端口，不显式指定也可以
|
||||
result = validate_url_safety("http://example.com:80/test", purpose="test")
|
||||
self.assertIn("example.com", result)
|
||||
|
||||
def test_port_443_allowed(self):
|
||||
"""端口 443 应该被允许."""
|
||||
result = validate_url_safety("https://example.com:443/test", purpose="test")
|
||||
self.assertIn("example.com", result)
|
||||
|
||||
def test_port_8080_rejected(self):
|
||||
"""非标准端口 8080 应该被拒绝."""
|
||||
with self.assertRaises(UrlSecurityError):
|
||||
validate_url_safety("http://example.com:8080/test", purpose="test")
|
||||
|
||||
def test_port_22_rejected(self):
|
||||
"""SSH 端口 22 应该被拒绝."""
|
||||
with self.assertRaises(UrlSecurityError):
|
||||
validate_url_safety("http://example.com:22/test", purpose="test")
|
||||
|
||||
# ── SSRF: 直接 IP 访问 ───────────────────────────────────────────────
|
||||
|
||||
def test_loopback_ip_rejected(self):
|
||||
"""回环地址 127.0.0.1 应该被拒绝."""
|
||||
with self.assertRaises(UrlSecurityError):
|
||||
validate_url_safety("http://127.0.0.1/test", purpose="test")
|
||||
|
||||
def test_private_ip_192_rejected(self):
|
||||
"""内网地址 192.168.x.x 应该被拒绝."""
|
||||
with self.assertRaises(UrlSecurityError):
|
||||
validate_url_safety("http://192.168.1.1/test", purpose="test")
|
||||
|
||||
def test_private_ip_10_rejected(self):
|
||||
"""内网地址 10.x.x.x 应该被拒绝."""
|
||||
with self.assertRaises(UrlSecurityError):
|
||||
validate_url_safety("http://10.0.0.1/test", purpose="test")
|
||||
|
||||
def test_private_ip_172_rejected(self):
|
||||
"""内网地址 172.16.x.x 应该被拒绝."""
|
||||
with self.assertRaises(UrlSecurityError):
|
||||
validate_url_safety("http://172.16.0.1/test", purpose="test")
|
||||
|
||||
def test_unspecified_ip_rejected(self):
|
||||
"""未指定地址 0.0.0.0 应该被拒绝."""
|
||||
with self.assertRaises(UrlSecurityError):
|
||||
validate_url_safety("http://0.0.0.0/test", purpose="test")
|
||||
|
||||
def test_ipv6_loopback_rejected(self):
|
||||
"""IPv6 回环 ::1 应该被拒绝."""
|
||||
with self.assertRaises(UrlSecurityError):
|
||||
validate_url_safety("http://[::1]/test", purpose="test")
|
||||
|
||||
def test_ipv6_link_local_rejected(self):
|
||||
"""IPv6 链路本地地址应该被拒绝."""
|
||||
with self.assertRaises(UrlSecurityError):
|
||||
validate_url_safety("http://[fe80::1]/test", purpose="test")
|
||||
|
||||
# ── SSRF: 内网主机名 ─────────────────────────────────────────────────
|
||||
|
||||
def test_localhost_rejected(self):
|
||||
"""localhost 应该被拒绝."""
|
||||
with self.assertRaises(UrlSecurityError):
|
||||
validate_url_safety("http://localhost/test", purpose="test")
|
||||
|
||||
def test_local_domain_rejected(self):
|
||||
""".local 域名应该被拒绝."""
|
||||
with self.assertRaises(UrlSecurityError):
|
||||
validate_url_safety("http://printer.local/test", purpose="test")
|
||||
|
||||
def test_internal_domain_rejected(self):
|
||||
""".internal 域名应该被拒绝."""
|
||||
with self.assertRaises(UrlSecurityError):
|
||||
validate_url_safety("http://db.internal/test", purpose="test")
|
||||
|
||||
# ── URL 格式校验 ─────────────────────────────────────────────────────
|
||||
|
||||
def test_empty_url_rejected(self):
|
||||
"""空 URL 应该被拒绝."""
|
||||
with self.assertRaises(UrlSecurityError):
|
||||
validate_url_safety("", purpose="test")
|
||||
|
||||
def test_none_url_rejected(self):
|
||||
"""None URL 应该被拒绝."""
|
||||
with self.assertRaises(UrlSecurityError):
|
||||
validate_url_safety(None, purpose="test") # type: ignore
|
||||
|
||||
def test_url_too_long_rejected(self):
|
||||
"""超长 URL 应该被拒绝."""
|
||||
long_url = "https://example.com/" + "a" * 3000
|
||||
with self.assertRaises(UrlSecurityError):
|
||||
validate_url_safety(long_url, purpose="test")
|
||||
|
||||
def test_no_hostname_rejected(self):
|
||||
"""缺少主机名应该被拒绝."""
|
||||
with self.assertRaises(UrlSecurityError):
|
||||
validate_url_safety("http:///test", purpose="test")
|
||||
|
||||
# ── is_url_safe 便捷函数 ─────────────────────────────────────────────
|
||||
|
||||
def test_is_url_safe_true(self):
|
||||
"""安全 URL 应该返回 True."""
|
||||
self.assertTrue(is_url_safe("https://example.com/test", purpose="test"))
|
||||
|
||||
def test_is_url_safe_false(self):
|
||||
"""不安全 URL 应该返回 False."""
|
||||
self.assertFalse(is_url_safe("http://127.0.0.1/test", purpose="test"))
|
||||
|
||||
def test_is_url_safe_empty(self):
|
||||
"""空 URL 应该返回 False."""
|
||||
self.assertFalse(is_url_safe("", purpose="test"))
|
||||
# ── validate_url_safety 基础校验 ─────────────────────────────────────────────
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
class TestValidateUrlSafetyBasics:
|
||||
"""URL 安全校验基础测试"""
|
||||
|
||||
def test_valid_http_url(self):
|
||||
url = "http://example.com/file.mp4"
|
||||
result = validate_url_safety(url)
|
||||
assert result == url
|
||||
|
||||
def test_valid_https_url(self):
|
||||
url = "https://example.com/file.mp4"
|
||||
result = validate_url_safety(url)
|
||||
assert result == url
|
||||
|
||||
def test_empty_url_raises(self):
|
||||
with pytest.raises(UrlSecurityError, match="为空"):
|
||||
validate_url_safety("")
|
||||
|
||||
def test_none_url_raises(self):
|
||||
with pytest.raises(UrlSecurityError):
|
||||
validate_url_safety(None)
|
||||
|
||||
def test_url_too_long_raises(self):
|
||||
long_url = "https://example.com/" + "a" * 2050
|
||||
with pytest.raises(UrlSecurityError, match="过长"):
|
||||
validate_url_safety(long_url)
|
||||
|
||||
def test_url_at_max_length_ok(self):
|
||||
base = "https://example.com/"
|
||||
pad = "a" * (MAX_URL_LENGTH - len(base))
|
||||
url = base + pad
|
||||
assert len(url) <= MAX_URL_LENGTH
|
||||
result = validate_url_safety(url)
|
||||
assert result == url
|
||||
|
||||
def test_invalid_scheme_ftp_raises(self):
|
||||
with pytest.raises(UrlSecurityError, match="scheme"):
|
||||
validate_url_safety("ftp://example.com/file")
|
||||
|
||||
def test_invalid_scheme_file_raises(self):
|
||||
with pytest.raises(UrlSecurityError, match="scheme"):
|
||||
validate_url_safety("file:///etc/passwd")
|
||||
|
||||
def test_invalid_scheme_data_raises(self):
|
||||
with pytest.raises(UrlSecurityError, match="scheme"):
|
||||
validate_url_safety("data:text/html,<script>")
|
||||
|
||||
def test_missing_scheme_raises(self):
|
||||
with pytest.raises(UrlSecurityError, match="scheme"):
|
||||
validate_url_safety("example.com/file")
|
||||
|
||||
def test_missing_hostname_raises(self):
|
||||
with pytest.raises(UrlSecurityError, match="主机名"):
|
||||
validate_url_safety("http:///path")
|
||||
|
||||
def test_uppercase_scheme_normalized(self):
|
||||
"""HTTP/HTTPS 大写也能通过"""
|
||||
url = "HTTPS://example.com/file"
|
||||
# scheme 检查用 lower 比较
|
||||
result = validate_url_safety(url)
|
||||
assert result == url
|
||||
|
||||
def test_default_port_80_ok(self):
|
||||
url = "http://example.com:80/file"
|
||||
result = validate_url_safety(url)
|
||||
assert result == url
|
||||
|
||||
def test_default_port_443_ok(self):
|
||||
url = "https://example.com:443/file"
|
||||
result = validate_url_safety(url)
|
||||
assert result == url
|
||||
|
||||
def test_non_standard_port_raises(self):
|
||||
with pytest.raises(UrlSecurityError, match="端口"):
|
||||
validate_url_safety("http://example.com:8080/file")
|
||||
|
||||
def test_port_22_ssh_raises(self):
|
||||
with pytest.raises(UrlSecurityError, match="端口"):
|
||||
validate_url_safety("http://example.com:22/file")
|
||||
|
||||
def test_port_3306_mysql_raises(self):
|
||||
with pytest.raises(UrlSecurityError, match="端口"):
|
||||
validate_url_safety("http://example.com:3306/file")
|
||||
|
||||
|
||||
class TestSafeDownload(unittest.TestCase):
|
||||
"""安全下载函数测试."""
|
||||
# ── 内网主机名 / SSRF 防护 ───────────────────────────────────────────────────
|
||||
|
||||
def setUp(self):
|
||||
self.temp_dir = tempfile.mkdtemp()
|
||||
|
||||
def tearDown(self):
|
||||
shutil.rmtree(self.temp_dir, ignore_errors=True)
|
||||
class TestInternalHostnameProtection:
|
||||
"""内网主机名防护测试"""
|
||||
|
||||
def test_safe_download_file_rejects_ssrf(self):
|
||||
"""SSRF 风险 URL 应该被拒绝下载."""
|
||||
dest = os.path.join(self.temp_dir, "test.bin")
|
||||
with self.assertRaises(UrlSecurityError):
|
||||
safe_download_file("http://127.0.0.1/test", dest, purpose="test")
|
||||
def test_localhost_raises(self):
|
||||
with pytest.raises(UrlSecurityError, match="内部主机名"):
|
||||
validate_url_safety("http://localhost/file")
|
||||
|
||||
def test_safe_download_bytes_rejects_ssrf(self):
|
||||
"""SSRF 风险 URL 应该被拒绝下载（bytes 版本）."""
|
||||
with self.assertRaises(UrlSecurityError):
|
||||
safe_download_bytes("http://localhost/test", purpose="test")
|
||||
def test_localhost_mixed_case_raises(self):
|
||||
with pytest.raises(UrlSecurityError):
|
||||
validate_url_safety("http://LocalHost/file")
|
||||
|
||||
def test_safe_download_file_size_limit(self):
|
||||
"""超过大小限制应该被拒绝."""
|
||||
# 用 mock server 测试太大的 content-length
|
||||
dest = os.path.join(self.temp_dir, "test.bin")
|
||||
# 直接验证参数：max_size=0 时任何下载都应超限
|
||||
# （这里用一个可访问的 URL 并设置极小的限制）
|
||||
# 为避免依赖外部网络，这里只测试函数参数传递
|
||||
with unittest.mock.patch("urllib.request.build_opener") as mock_opener:
|
||||
mock_resp = unittest.mock.MagicMock()
|
||||
mock_resp.headers = {"Content-Length": "1000"}
|
||||
mock_resp.read.return_value = b""
|
||||
mock_opener.return_value.open.return_value = mock_resp
|
||||
# 设置 max_size=500，content-length=1000 应被拒绝
|
||||
with self.assertRaises(UrlSecurityError):
|
||||
safe_download_file(
|
||||
"https://example.com/test",
|
||||
def test_localhost_localdomain_raises(self):
|
||||
with pytest.raises(UrlSecurityError):
|
||||
_check_internal_hostnames("localhost.localdomain")
|
||||
|
||||
def test_metadata_hostname_raises(self):
|
||||
with pytest.raises(UrlSecurityError):
|
||||
_check_internal_hostnames("metadata")
|
||||
|
||||
def test_metadata_google_internal_raises(self):
|
||||
with pytest.raises(UrlSecurityError):
|
||||
_check_internal_hostnames("metadata.google.internal")
|
||||
|
||||
def test_dot_local_domain_raises(self):
|
||||
with pytest.raises(UrlSecurityError, match="内网域名"):
|
||||
validate_url_safety("http://myservice.local/file")
|
||||
|
||||
def test_dot_internal_domain_raises(self):
|
||||
with pytest.raises(UrlSecurityError, match="内网域名"):
|
||||
validate_url_safety("http://myservice.internal/file")
|
||||
|
||||
def test_dot_localdomain_raises(self):
|
||||
with pytest.raises(UrlSecurityError):
|
||||
_check_internal_hostnames("server.localdomain")
|
||||
|
||||
def test_loopback_ip_127_0_0_1_raises(self):
|
||||
with pytest.raises(UrlSecurityError, match="直接 IP|回环"):
|
||||
validate_url_safety("http://127.0.0.1/file")
|
||||
|
||||
def test_metadata_ip_169_254_raises(self):
|
||||
"""云元数据服务 IP"""
|
||||
with pytest.raises(UrlSecurityError):
|
||||
validate_url_safety("http://169.254.169.254/latest/meta-data/")
|
||||
|
||||
def test_private_ip_10_raises(self):
|
||||
with pytest.raises(UrlSecurityError):
|
||||
validate_url_safety("http://10.0.0.1/file")
|
||||
|
||||
def test_private_ip_172_16_raises(self):
|
||||
with pytest.raises(UrlSecurityError):
|
||||
validate_url_safety("http://172.16.0.1/file")
|
||||
|
||||
def test_private_ip_192_168_raises(self):
|
||||
with pytest.raises(UrlSecurityError):
|
||||
validate_url_safety("http://192.168.1.1/file")
|
||||
|
||||
def test_unspecified_ip_0_0_0_0_raises(self):
|
||||
with pytest.raises(UrlSecurityError):
|
||||
validate_url_safety("http://0.0.0.0/file")
|
||||
|
||||
def test_ipv6_loopback_raises(self):
|
||||
with pytest.raises(UrlSecurityError):
|
||||
validate_url_safety("http://[::1]/file")
|
||||
|
||||
def test_public_ip_ok(self):
|
||||
"""公网IP在ALLOW_DIRECT_IP默认关闭时应被拦截"""
|
||||
# 默认 ALLOW_DIRECT_IP = false
|
||||
with pytest.raises(UrlSecurityError, match="直接 IP"):
|
||||
validate_url_safety("http://8.8.8.8/file")
|
||||
|
||||
|
||||
# ── 可信域名白名单 ───────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestTrustedDomains:
|
||||
"""可信域名白名单测试"""
|
||||
|
||||
def test_is_trusted_domain_exact_match(self):
|
||||
with patch("packages.shared.url_security.TRUSTED_DOMAINS", {"example.com", "cdn.example.org"}):
|
||||
# 重新加载模块以应用环境变量不太现实，直接测函数
|
||||
# 直接改全局状态再还原
|
||||
import packages.shared.url_security as mod
|
||||
from packages.shared.url_security import _is_trusted_domain
|
||||
|
||||
original = mod.TRUSTED_DOMAINS
|
||||
mod.TRUSTED_DOMAINS = {"example.com", "cdn.example.org"}
|
||||
try:
|
||||
assert _is_trusted_domain("example.com") is True
|
||||
assert _is_trusted_domain("cdn.example.org") is True
|
||||
finally:
|
||||
mod.TRUSTED_DOMAINS = original
|
||||
|
||||
def test_is_trusted_domain_subdomain(self):
|
||||
import packages.shared.url_security as mod
|
||||
|
||||
original = mod.TRUSTED_DOMAINS
|
||||
mod.TRUSTED_DOMAINS = {"example.com"}
|
||||
try:
|
||||
assert mod._is_trusted_domain("sub.example.com") is True
|
||||
assert mod._is_trusted_domain("a.b.example.com") is True
|
||||
finally:
|
||||
mod.TRUSTED_DOMAINS = original
|
||||
|
||||
def test_is_trusted_domain_no_match(self):
|
||||
import packages.shared.url_security as mod
|
||||
|
||||
original = mod.TRUSTED_DOMAINS
|
||||
mod.TRUSTED_DOMAINS = {"example.com"}
|
||||
try:
|
||||
assert mod._is_trusted_domain("other.com") is False
|
||||
assert mod._is_trusted_domain("notexample.com") is False
|
||||
finally:
|
||||
mod.TRUSTED_DOMAINS = original
|
||||
|
||||
def test_validate_with_trusted_domains_restricted(self):
|
||||
"""白名单非空时，不在白名单中的域名被拒"""
|
||||
import packages.shared.url_security as mod
|
||||
|
||||
original = mod.TRUSTED_DOMAINS
|
||||
mod.TRUSTED_DOMAINS = {"trusted.com"}
|
||||
try:
|
||||
# 不在白名单中 - 在 _is_trusted_domain 检查时就被拒，不走 DNS
|
||||
with pytest.raises(UrlSecurityError, match="白名单"):
|
||||
validate_url_safety("https://untrusted.com/file")
|
||||
|
||||
# 在白名单中 - 需要 mock DNS 解析避免实际网络请求
|
||||
with patch("packages.shared.url_security._check_ssrf_domain"):
|
||||
result = validate_url_safety("https://trusted.com/file")
|
||||
assert result == "https://trusted.com/file"
|
||||
# 子域名
|
||||
result = validate_url_safety("https://sub.trusted.com/file")
|
||||
assert result == "https://sub.trusted.com/file"
|
||||
finally:
|
||||
mod.TRUSTED_DOMAINS = original
|
||||
|
||||
|
||||
# ── is_url_safe 便捷函数 ─────────────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestIsUrlSafe:
|
||||
"""is_url_safe 便捷函数测试"""
|
||||
|
||||
def test_safe_url_returns_true(self):
|
||||
assert is_url_safe("https://example.com/file") is True
|
||||
|
||||
def test_unsafe_url_returns_false(self):
|
||||
assert is_url_safe("http://localhost/file") is False
|
||||
|
||||
def test_empty_url_returns_false(self):
|
||||
assert is_url_safe("") is False
|
||||
|
||||
def test_invalid_scheme_returns_false(self):
|
||||
assert is_url_safe("ftp://example.com/file") is False
|
||||
|
||||
|
||||
# ── UrlSecurityError 异常类 ───────────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestUrlSecurityError:
|
||||
"""UrlSecurityError 异常类测试"""
|
||||
|
||||
def test_is_value_error_subclass(self):
|
||||
assert issubclass(UrlSecurityError, ValueError)
|
||||
|
||||
def test_error_message(self):
|
||||
err = UrlSecurityError("test message")
|
||||
assert str(err) == "test message"
|
||||
|
||||
|
||||
# ── NoRedirectHandler ────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestNoRedirectHandler:
|
||||
"""NoRedirectHandler 测试"""
|
||||
|
||||
def test_redirect_request_returns_none(self):
|
||||
handler = NoRedirectHandler()
|
||||
result = handler.redirect_request(
|
||||
MagicMock(),
|
||||
MagicMock(),
|
||||
302,
|
||||
"Found",
|
||||
{"Location": "http://other.com"},
|
||||
"http://other.com",
|
||||
)
|
||||
assert result is None
|
||||
|
||||
|
||||
# ── 魔数校验 ─────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestMagicNumberValidation:
|
||||
"""文件魔数校验测试"""
|
||||
|
||||
def test_valid_png(self, tmp_path):
|
||||
f = tmp_path / "test.png"
|
||||
# PNG 文件头: 89 50 4E 47 0D 0A 1A 0A
|
||||
f.write_bytes(b"\x89PNG\r\n\x1a\n" + b"\x00" * 100)
|
||||
# 不抛异常 = 通过
|
||||
_validate_magic_number(str(f), ALLOWED_IMAGE_MIME_TYPES)
|
||||
|
||||
def test_valid_jpeg(self, tmp_path):
|
||||
f = tmp_path / "test.jpg"
|
||||
# JPEG 文件头: FF D8 FF
|
||||
f.write_bytes(b"\xff\xd8\xff\xe0" + b"\x00" * 100)
|
||||
_validate_magic_number(str(f), ALLOWED_IMAGE_MIME_TYPES)
|
||||
|
||||
def test_valid_gif87a(self, tmp_path):
|
||||
f = tmp_path / "test.gif"
|
||||
f.write_bytes(b"GIF87a" + b"\x00" * 100)
|
||||
_validate_magic_number(str(f), ALLOWED_IMAGE_MIME_TYPES)
|
||||
|
||||
def test_valid_gif89a(self, tmp_path):
|
||||
f = tmp_path / "test.gif"
|
||||
f.write_bytes(b"GIF89a" + b"\x00" * 100)
|
||||
_validate_magic_number(str(f), ALLOWED_IMAGE_MIME_TYPES)
|
||||
|
||||
def test_valid_webp(self, tmp_path):
|
||||
f = tmp_path / "test.webp"
|
||||
# RIFF....WEBP
|
||||
data = bytearray(b"RIFF")
|
||||
data += b"\x00\x00\x00\x00" # size placeholder
|
||||
data += b"WEBP"
|
||||
data += b"\x00" * 100
|
||||
f.write_bytes(bytes(data))
|
||||
_validate_magic_number(str(f), ALLOWED_IMAGE_MIME_TYPES)
|
||||
|
||||
def test_valid_bmp(self, tmp_path):
|
||||
f = tmp_path / "test.bmp"
|
||||
f.write_bytes(b"BM" + b"\x00" * 100)
|
||||
_validate_magic_number(str(f), ALLOWED_IMAGE_MIME_TYPES)
|
||||
|
||||
def test_valid_wav(self, tmp_path):
|
||||
f = tmp_path / "test.wav"
|
||||
# RIFF....WAVE
|
||||
data = bytearray(b"RIFF")
|
||||
data += b"\x00\x00\x00\x00"
|
||||
data += b"WAVE"
|
||||
data += b"\x00" * 100
|
||||
f.write_bytes(bytes(data))
|
||||
_validate_magic_number(str(f), ALLOWED_AUDIO_MIME_TYPES)
|
||||
|
||||
def test_valid_mp3_id3(self, tmp_path):
|
||||
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), ALLOWED_AUDIO_MIME_TYPES)
|
||||
|
||||
def test_valid_mp3_adts(self, tmp_path):
|
||||
f = tmp_path / "test.mp3"
|
||||
f.write_bytes(b"\xff\xfb\x90\x00" + b"\x00" * 100)
|
||||
_validate_magic_number(str(f), ALLOWED_AUDIO_MIME_TYPES)
|
||||
|
||||
def test_valid_ogg(self, tmp_path):
|
||||
f = tmp_path / "test.ogg"
|
||||
f.write_bytes(b"OggS\x00\x02\x00\x00" + b"\x00" * 100)
|
||||
_validate_magic_number(str(f), ALLOWED_AUDIO_MIME_TYPES)
|
||||
|
||||
def test_valid_flac(self, tmp_path):
|
||||
f = tmp_path / "test.flac"
|
||||
f.write_bytes(b"fLaC" + b"\x00" * 100)
|
||||
_validate_magic_number(str(f), ALLOWED_AUDIO_MIME_TYPES)
|
||||
|
||||
def test_invalid_file_content_raises(self, tmp_path):
|
||||
f = tmp_path / "test.bin"
|
||||
f.write_bytes(b"this is not an image file at all")
|
||||
with pytest.raises(UrlSecurityError, match="魔数"):
|
||||
_validate_magic_number(str(f), ALLOWED_IMAGE_MIME_TYPES)
|
||||
|
||||
def test_empty_file_raises(self, tmp_path):
|
||||
f = tmp_path / "empty.bin"
|
||||
f.write_bytes(b"")
|
||||
with pytest.raises(UrlSecurityError, match="为空"):
|
||||
_validate_magic_number(str(f), ALLOWED_IMAGE_MIME_TYPES)
|
||||
|
||||
def test_nonexistent_file_raises(self, tmp_path):
|
||||
with pytest.raises(UrlSecurityError, match="读取文件头失败"):
|
||||
_validate_magic_number(str(tmp_path / "no_such_file"), ALLOWED_IMAGE_MIME_TYPES)
|
||||
|
||||
def test_no_allowed_mime_types_skips(self, tmp_path):
|
||||
"""allowed_mime_types 为空时跳过校验"""
|
||||
f = tmp_path / "test.bin"
|
||||
f.write_bytes(b"random data here")
|
||||
# 不抛异常
|
||||
_validate_magic_number(str(f), set())
|
||||
|
||||
def test_unknown_mime_types_skips(self, tmp_path):
|
||||
"""没有已知魔数的 MIME 类型跳过校验"""
|
||||
f = tmp_path / "test.bin"
|
||||
f.write_bytes(b"random data")
|
||||
_validate_magic_number(str(f), {"application/x-unknown-type"})
|
||||
|
||||
|
||||
# ── safe_download_file (mock 网络) ───────────────────────────────────────────
|
||||
|
||||
|
||||
class TestSafeDownloadFile:
|
||||
"""safe_download_file 下载测试（mock 网络）"""
|
||||
|
||||
def test_download_success(self, tmp_path):
|
||||
test_content = b"Hello, this is test file content!"
|
||||
dest = str(tmp_path / "output.bin")
|
||||
|
||||
mock_resp = MagicMock()
|
||||
mock_resp.headers = {"Content-Type": "application/octet-stream"}
|
||||
mock_resp.read.side_effect = [test_content, b""]
|
||||
|
||||
with patch("packages.shared.url_security.NoRedirectHandler") as mock_handler_cls:
|
||||
mock_handler = MagicMock()
|
||||
mock_handler_cls.return_value = mock_handler
|
||||
|
||||
mock_opener = MagicMock()
|
||||
mock_opener.open.return_value = mock_resp
|
||||
|
||||
with patch("urllib.request.build_opener", return_value=mock_opener):
|
||||
size = safe_download_file(
|
||||
"https://example.com/test.bin",
|
||||
dest,
|
||||
purpose="test",
|
||||
max_size=500,
|
||||
)
|
||||
|
||||
def test_safe_download_file_mime_rejected(self):
|
||||
"""不允许的 MIME 类型应该被拒绝."""
|
||||
dest = os.path.join(self.temp_dir, "test.bin")
|
||||
with unittest.mock.patch("urllib.request.build_opener") as mock_opener:
|
||||
mock_resp = unittest.mock.MagicMock()
|
||||
mock_resp.headers = {"Content-Type": "text/html"}
|
||||
mock_resp.read.return_value = b""
|
||||
mock_opener.return_value.open.return_value = mock_resp
|
||||
with self.assertRaises(UrlSecurityError):
|
||||
safe_download_file(
|
||||
"https://example.com/test.mp3",
|
||||
dest,
|
||||
purpose="test",
|
||||
allowed_mime_types=ALLOWED_AUDIO_MIME_TYPES,
|
||||
)
|
||||
assert size == len(test_content)
|
||||
with open(dest, "rb") as f:
|
||||
assert f.read() == test_content
|
||||
|
||||
def test_download_with_mime_check_passes(self, tmp_path):
|
||||
# PNG 文件
|
||||
test_content = b"\x89PNG\r\n\x1a\n" + b"\x00" * 200
|
||||
dest = str(tmp_path / "test.png")
|
||||
|
||||
mock_resp = MagicMock()
|
||||
mock_resp.headers = {"Content-Type": "image/png"}
|
||||
mock_resp.read.side_effect = [test_content, b""]
|
||||
|
||||
with patch("urllib.request.build_opener") as mock_build:
|
||||
mock_opener = MagicMock()
|
||||
mock_opener.open.return_value = mock_resp
|
||||
mock_build.return_value = mock_opener
|
||||
|
||||
def test_safe_download_file_mime_allowed(self):
|
||||
"""允许的 MIME 类型应该通过."""
|
||||
dest = os.path.join(self.temp_dir, "test.mp3")
|
||||
with unittest.mock.patch("urllib.request.build_opener") as mock_opener:
|
||||
mock_resp = unittest.mock.MagicMock()
|
||||
mock_resp.headers = {"Content-Type": "audio/mpeg"}
|
||||
mock_resp.read.side_effect = [b"ID3audio_data", b""]
|
||||
mock_resp.geturl.return_value = "https://example.com/test.mp3"
|
||||
mock_opener.return_value.open.return_value = mock_resp
|
||||
size = safe_download_file(
|
||||
"https://example.com/test.mp3",
|
||||
"https://example.com/test.png",
|
||||
dest,
|
||||
purpose="test",
|
||||
allowed_mime_types=ALLOWED_AUDIO_MIME_TYPES,
|
||||
allowed_mime_types={"image/png", "image/jpeg"},
|
||||
)
|
||||
self.assertEqual(size, 13)
|
||||
self.assertTrue(os.path.exists(dest))
|
||||
|
||||
def test_safe_download_file_stream_size_limit(self):
|
||||
"""流式下载时超过大小限制应该中断."""
|
||||
dest = os.path.join(self.temp_dir, "test.bin")
|
||||
with unittest.mock.patch("urllib.request.build_opener") as mock_opener:
|
||||
mock_resp = unittest.mock.MagicMock()
|
||||
mock_resp.headers = {}
|
||||
# 每次返回 100 字节，max_size=500，第 6 次读取就超限
|
||||
mock_resp.read.side_effect = lambda n: b"x" * n if n < 1000 else b"x" * 100
|
||||
# 改成返回固定 100 字节，直到第 N 次后返回空
|
||||
call_count = [0]
|
||||
assert size == len(test_content)
|
||||
|
||||
def mock_read(size):
|
||||
call_count[0] += 1
|
||||
if call_count[0] > 10:
|
||||
return b""
|
||||
return b"x" * 100
|
||||
def test_download_mime_type_rejected(self, tmp_path):
|
||||
test_content = b"GIF89a" + b"\x00" * 50
|
||||
dest = str(tmp_path / "test.gif")
|
||||
|
||||
mock_resp.read = mock_read
|
||||
mock_opener.return_value.open.return_value = mock_resp
|
||||
with self.assertRaises(UrlSecurityError):
|
||||
mock_resp = MagicMock()
|
||||
mock_resp.headers = {"Content-Type": "image/gif"}
|
||||
mock_resp.read.side_effect = [test_content, b""]
|
||||
|
||||
with patch("urllib.request.build_opener") as mock_build:
|
||||
mock_opener = MagicMock()
|
||||
mock_opener.open.return_value = mock_resp
|
||||
mock_build.return_value = mock_opener
|
||||
|
||||
with pytest.raises(UrlSecurityError, match="Content-Type"):
|
||||
safe_download_file(
|
||||
"https://example.com/test",
|
||||
"https://example.com/test.gif",
|
||||
dest,
|
||||
purpose="test",
|
||||
max_size=500, # 500 字节上限
|
||||
allowed_mime_types={"image/png"},
|
||||
)
|
||||
|
||||
def test_safe_download_bytes_returns_content(self):
|
||||
"""safe_download_bytes 应该返回文件内容."""
|
||||
test_data = b"ID3hello world test audio"
|
||||
with unittest.mock.patch("urllib.request.build_opener") as mock_opener:
|
||||
mock_resp = unittest.mock.MagicMock()
|
||||
mock_resp.headers = {"Content-Type": "audio/mpeg"}
|
||||
call_count = [0]
|
||||
def test_download_size_limit_exceeded(self, tmp_path):
|
||||
"""流式下载时超过大小限制被中断（无 Content-Length header）"""
|
||||
dest = str(tmp_path / "big.bin")
|
||||
chunk = b"x" * 1024 # 1KB chunks
|
||||
|
||||
def mock_read(size):
|
||||
call_count[0] += 1
|
||||
if call_count[0] > 1:
|
||||
return b""
|
||||
return test_data
|
||||
mock_resp = MagicMock()
|
||||
# 没有 Content-Length header，走流式检查
|
||||
mock_resp.headers = {"Content-Type": "application/octet-stream"}
|
||||
# 模拟多次读取，超过 5KB 限制（6个chunk = 6KB）
|
||||
mock_resp.read.side_effect = [chunk] * 6 + [b""]
|
||||
|
||||
with patch("urllib.request.build_opener") as mock_build:
|
||||
mock_opener = MagicMock()
|
||||
mock_opener.open.return_value = mock_resp
|
||||
mock_build.return_value = mock_opener
|
||||
|
||||
with pytest.raises(UrlSecurityError, match="超过大小限制"):
|
||||
safe_download_file(
|
||||
"https://example.com/big.bin",
|
||||
dest,
|
||||
purpose="test",
|
||||
max_size=5000, # 5KB limit
|
||||
)
|
||||
|
||||
def test_download_content_length_too_large(self, tmp_path):
|
||||
dest = str(tmp_path / "big.bin")
|
||||
|
||||
mock_resp = MagicMock()
|
||||
mock_resp.headers = {"Content-Type": "application/octet-stream", "Content-Length": "1000000"}
|
||||
mock_resp.read.side_effect = [b"data"]
|
||||
|
||||
with patch("urllib.request.build_opener") as mock_build:
|
||||
mock_opener = MagicMock()
|
||||
mock_opener.open.return_value = mock_resp
|
||||
mock_build.return_value = mock_opener
|
||||
|
||||
with pytest.raises(UrlSecurityError, match="文件过大"):
|
||||
safe_download_file(
|
||||
"https://example.com/big.bin",
|
||||
dest,
|
||||
purpose="test",
|
||||
max_size=500000,
|
||||
)
|
||||
|
||||
def test_download_localhost_rejected(self, tmp_path):
|
||||
"""内网 URL 在下载前就被拒"""
|
||||
dest = str(tmp_path / "out.bin")
|
||||
with pytest.raises(UrlSecurityError):
|
||||
safe_download_file("http://localhost/file", dest)
|
||||
|
||||
def test_download_invalid_scheme_rejected(self, tmp_path):
|
||||
dest = str(tmp_path / "out.bin")
|
||||
with pytest.raises(UrlSecurityError, match="scheme"):
|
||||
safe_download_file("ftp://example.com/file", dest)
|
||||
|
||||
|
||||
# ── safe_download_bytes ───────────────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestSafeDownloadBytes:
|
||||
"""safe_download_bytes 测试"""
|
||||
|
||||
def test_download_returns_bytes(self, tmp_path):
|
||||
test_content = b"hello bytes download test"
|
||||
|
||||
mock_resp = MagicMock()
|
||||
mock_resp.headers = {"Content-Type": "application/octet-stream"}
|
||||
mock_resp.read.side_effect = [test_content, b""]
|
||||
|
||||
with patch("urllib.request.build_opener") as mock_build:
|
||||
mock_opener = MagicMock()
|
||||
mock_opener.open.return_value = mock_resp
|
||||
mock_build.return_value = mock_opener
|
||||
|
||||
mock_resp.read = mock_read
|
||||
mock_opener.return_value.open.return_value = mock_resp
|
||||
result = safe_download_bytes(
|
||||
"https://example.com/test.mp3",
|
||||
"https://example.com/test.bin",
|
||||
purpose="test",
|
||||
allowed_mime_types=ALLOWED_AUDIO_MIME_TYPES,
|
||||
)
|
||||
self.assertEqual(result, test_data)
|
||||
|
||||
assert result == test_content
|
||||
assert isinstance(result, bytes)
|
||||
|
||||
def test_download_unsafe_url_raises(self):
|
||||
with pytest.raises(UrlSecurityError):
|
||||
safe_download_bytes("http://127.0.0.1/secret")
|
||||
|
||||
|
||||
# ── 常量导出验证 ─────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestConstants:
|
||||
"""模块常量验证"""
|
||||
|
||||
def test_allowed_schemes(self):
|
||||
assert "http" in ALLOWED_SCHEMES
|
||||
assert "https" in ALLOWED_SCHEMES
|
||||
assert len(ALLOWED_SCHEMES) == 2
|
||||
|
||||
def test_allowed_ports(self):
|
||||
assert 80 in ALLOWED_PORTS
|
||||
assert 443 in ALLOWED_PORTS
|
||||
assert len(ALLOWED_PORTS) == 2
|
||||
|
||||
def test_max_url_length(self):
|
||||
assert MAX_URL_LENGTH == 2048
|
||||
|
||||
Reference in New Issue
Block a user