xiaoxia-saas/tests/unit/test_url_security.py

"""
url_security URL安全校验单元测试

覆盖：
- validate_url_safety: scheme/主机/端口/SSRF/内网域名/白名单
- is_url_safe: 便捷函数
- UrlSecurityError / NoRedirectHandler
- _validate_magic_number: 文件魔数校验
- safe_download_file / safe_download_bytes: mock 网络测试
"""

import ipaddress
import os
import tempfile
from unittest.mock import MagicMock, patch

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,
    _check_ssrf_ip,
    _is_trusted_domain,
    _validate_magic_number,
    is_url_safe,
    safe_download_bytes,
    safe_download_file,
    validate_url_safety,
)

# ── validate_url_safety 基础校验 ─────────────────────────────────────────────


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")


# ── 内网主机名 / SSRF 防护 ───────────────────────────────────────────────────


class TestInternalHostnameProtection:
    """内网主机名防护测试"""

    def test_localhost_raises(self):
        with pytest.raises(UrlSecurityError, match="内部主机名"):
            validate_url_safety("http://localhost/file")

    def test_localhost_mixed_case_raises(self):
        with pytest.raises(UrlSecurityError):
            validate_url_safety("http://LocalHost/file")

    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",
                )

        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

            size = safe_download_file(
                "https://example.com/test.png",
                dest,
                purpose="test",
                allowed_mime_types={"image/png", "image/jpeg"},
            )

        assert size == len(test_content)

    def test_download_mime_type_rejected(self, tmp_path):
        test_content = b"GIF89a" + b"\x00" * 50
        dest = str(tmp_path / "test.gif")

        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.gif",
                    dest,
                    purpose="test",
                    allowed_mime_types={"image/png"},
                )

    def test_download_size_limit_exceeded(self, tmp_path):
        """流式下载时超过大小限制被中断（无 Content-Length header）"""
        dest = str(tmp_path / "big.bin")
        chunk = b"x" * 1024  # 1KB chunks

        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

            result = safe_download_bytes(
                "https://example.com/test.bin",
                purpose="test",
            )

        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


# ── SSRF IP 检查详细覆盖 ──────────────────────────────────────────────────────


class TestSSRFIPCheck:
    """_check_ssrf_ip 各类型 IP 拦截覆盖."""

    def test_loopback_ipv4_blocked(self):
        """IPv4 回环 127.0.0.1 被拦."""
        with pytest.raises(UrlSecurityError):
            _check_ssrf_ip(ipaddress.ip_address("127.0.0.1"))

    def test_loopback_ipv4_another_blocked(self):
        """127.x 其他段也被拦."""
        with pytest.raises(UrlSecurityError):
            _check_ssrf_ip(ipaddress.ip_address("127.255.255.1"))

    def test_loopback_ipv6_blocked(self):
        """IPv6 回环 ::1 被拦."""
        with pytest.raises(UrlSecurityError):
            _check_ssrf_ip(ipaddress.ip_address("::1"))

    def test_private_10_range_blocked(self):
        """10.0.0.0/8 私有段被拦."""
        with pytest.raises(UrlSecurityError):
            _check_ssrf_ip(ipaddress.ip_address("10.255.255.255"))

    def test_private_172_range_blocked(self):
        """172.16.0.0/12 私有段被拦."""
        with pytest.raises(UrlSecurityError):
            _check_ssrf_ip(ipaddress.ip_address("172.31.255.255"))

    def test_private_192_range_blocked(self):
        """192.168.0.0/16 私有段被拦."""
        with pytest.raises(UrlSecurityError):
            _check_ssrf_ip(ipaddress.ip_address("192.168.255.255"))

    def test_link_local_ipv4_blocked(self):
        """169.254.x.x 链路本地被拦."""
        with pytest.raises(UrlSecurityError):
            _check_ssrf_ip(ipaddress.ip_address("169.254.1.1"))

    def test_multicast_ipv4_blocked(self):
        """224.x 组播被拦."""
        with pytest.raises(UrlSecurityError):
            _check_ssrf_ip(ipaddress.ip_address("224.0.0.251"))

    def test_unspecified_ipv4_blocked(self):
        """0.0.0.0 未指定地址被拦."""
        with pytest.raises(UrlSecurityError):
            _check_ssrf_ip(ipaddress.ip_address("0.0.0.0"))

    def test_unspecified_ipv6_blocked(self):
        """:: 未指定地址被拦."""
        with pytest.raises(UrlSecurityError):
            _check_ssrf_ip(ipaddress.ip_address("::"))

    def test_reserved_ipv4_blocked(self):
        """240.0.0.0/4 保留段被拦（含在 is_private 或 is_reserved 中）."""
        with pytest.raises(UrlSecurityError):
            _check_ssrf_ip(ipaddress.ip_address("240.0.0.1"))

    def test_public_ipv4_passes(self):
        """公网 IPv4 通过 _check_ssrf_ip."""
        _check_ssrf_ip(ipaddress.ip_address("8.8.8.8"))

    def test_public_ipv4_another_passes(self):
        """另一个公网 IPv4 通过."""
        _check_ssrf_ip(ipaddress.ip_address("1.1.1.1"))

    def test_public_ipv6_passes(self):
        """公网 IPv6 通过."""
        _check_ssrf_ip(ipaddress.ip_address("2001:4860:4860::8888"))


# ── 直接 IP 访问拦截 ────────────────────────────────────────────────────────


class TestDirectIPAccess:
    """直接 IP 访问控制（ALLOW_DIRECT_IP 开关）."""

    def test_direct_ipv4_blocked_by_default(self):
        """默认禁止直接 IP 访问."""
        with pytest.raises(UrlSecurityError, match="禁止直接 IP 访问"):
            validate_url_safety("https://8.8.8.8/audio.mp3")

    def test_direct_private_ip_blocked_even_with_flag(self):
        """ALLOW_DIRECT_IP=true 时私有 IP 仍被 SSRF 拦."""
        with patch("packages.shared.url_security.ALLOW_DIRECT_IP", True):
            with pytest.raises(UrlSecurityError):
                validate_url_safety("https://192.168.1.1/a.mp3")

    def test_direct_ip_allowed_when_flag_on(self):
        """ALLOW_DIRECT_IP=true 时公网 IP 通过."""
        with patch("packages.shared.url_security.ALLOW_DIRECT_IP", True):
            # 用 mock 绕过 DNS 解析路径，走 IP 分支
            with patch("packages.shared.url_security._check_ssrf_ip") as mock_check:
                result = validate_url_safety("https://8.8.8.8/a.mp3")
                assert "8.8.8.8" in result
                mock_check.assert_called_once()


# ── 魔数校验补充覆盖 ────────────────────────────────────────────────────────


class TestMagicNumberExtended:
    """魔数校验补充：更多格式 + 边界场景."""

    def test_aac_adts_mpeg4(self):
        """AAC ADTS MPEG-4 魔数通过."""
        with tempfile.NamedTemporaryFile(delete=False) as f:
            f.write(b"\xff\xf1" + b"\x00" * 50)
            tmp = f.name
        try:
            _validate_magic_number(tmp, {"audio/aac"})
        finally:
            os.unlink(tmp)

    def test_m4a_ftyp_magic(self):
        """M4A ftyp 魔数通过."""
        with tempfile.NamedTemporaryFile(delete=False) as f:
            f.write(b"\x00\x00\x00\x20ftypM4A " + b"\x00" * 50)
            tmp = f.name
        try:
            _validate_magic_number(tmp, {"audio/x-m4a"})
        finally:
            os.unlink(tmp)

    def test_webp_riff_webp(self):
        """WebP RIFF+WEBP 魔数通过."""
        with tempfile.NamedTemporaryFile(delete=False) as f:
            f.write(b"RIFF\x00\x00\x00\x00WEBP" + b"\x00" * 50)
            tmp = f.name
        try:
            _validate_magic_number(tmp, {"image/webp"})
        finally:
            os.unlink(tmp)

    def test_video_mp4_magic(self):
        """video/mp4 ftyp 魔数通过."""
        with tempfile.NamedTemporaryFile(delete=False) as f:
            f.write(b"\x00\x00\x00\x20ftypmp42" + b"\x00" * 50)
            tmp = f.name
        try:
            _validate_magic_number(tmp, {"video/mp4"})
        finally:
            os.unlink(tmp)

    def test_matroska_webm_magic(self):
        """WebM EBML 魔数通过."""
        with tempfile.NamedTemporaryFile(delete=False) as f:
            f.write(b"\x1a\x45\xdf\xa3" + b"\x00" * 50)
            tmp = f.name
        try:
            _validate_magic_number(tmp, {"video/webm"})
        finally:
            os.unlink(tmp)

    def test_bmp_magic(self):
        """BMP BM 魔数通过."""
        with tempfile.NamedTemporaryFile(delete=False) as f:
            f.write(b"BM" + b"\x00" * 50)
            tmp = f.name
        try:
            _validate_magic_number(tmp, {"image/bmp"})
        finally:
            os.unlink(tmp)

    def test_flac_magic(self):
        """FLAC fLaC 魔数通过."""
        with tempfile.NamedTemporaryFile(delete=False) as f:
            f.write(b"fLaC" + b"\x00" * 50)
            tmp = f.name
        try:
            _validate_magic_number(tmp, {"audio/flac"})
        finally:
            os.unlink(tmp)

    def test_ogg_magic(self):
        """OGG OggS 魔数通过."""
        with tempfile.NamedTemporaryFile(delete=False) as f:
            f.write(b"OggS" + b"\x00" * 50)
            tmp = f.name
        try:
            _validate_magic_number(tmp, {"audio/ogg"})
        finally:
            os.unlink(tmp)

    def test_read_error_raises_security_error(self):
        """文件读取失败包装为 UrlSecurityError."""
        with pytest.raises(UrlSecurityError, match="读取文件头失败"):
            _validate_magic_number("/nonexistent/path/file.mp3", {"audio/mpeg"})

    def test_multi_type_one_match(self):
        """多类型白名单，只要一个匹配就通过."""
        with tempfile.NamedTemporaryFile(delete=False) as f:
            f.write(b"\xff\xd8\xff" + b"\x00" * 50)
            tmp = f.name
        try:
            _validate_magic_number(tmp, {"image/png", "image/jpeg", "image/gif"})
        finally:
            os.unlink(tmp)

    def test_application_octet_stream_skipped(self):
        """application/octet-stream 没有专属魔数，跳过校验."""
        # 注意：octet-stream 在 _MAGIC_NUMBERS 中没有条目，所以跳过
        # 但实际白名单中常包含它，所以它的存在不应阻断
        with tempfile.NamedTemporaryFile(delete=False) as f:
            f.write(b"random stuff")
            tmp = f.name
        try:
            # octet-stream + png，png有魔数，png不匹配就会失败
            # 只有 octet-stream 时应该跳过
            _validate_magic_number(tmp, {"application/octet-stream"})
        finally:
            os.unlink(tmp)