xiaoxia-saas/tests/unit/test_url_security.py

"""URL安全校验纯逻辑测试 — SSRF防护/魔数校验/可信域名.

聚焦纯函数校验逻辑，DNS解析和实际下载用mock隔离。
"""

from __future__ import annotations

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

import pytest

from shared.url_security import (
    ALLOWED_PORTS,
    ALLOWED_SCHEMES,
    UrlSecurityError,
    _check_internal_hostnames,
    _check_ssrf_ip,
    _is_trusted_domain,
    _validate_magic_number,
    is_url_safe,
    validate_url_safety,
)


class TestUrlSecurityError:
    """UrlSecurityError 异常类."""

    def test_is_value_error(self):
        """继承自 ValueError."""
        assert issubclass(UrlSecurityError, ValueError)

    def test_message_preserved(self):
        """错误消息被保留."""
        err = UrlSecurityError("test message")
        assert str(err) == "test message"


class TestValidateUrlSafetyBasic:
    """validate_url_safety 基础校验（不涉及DNS解析）."""

    def test_empty_url_rejected(self):
        """空URL被拒绝."""
        with pytest.raises(UrlSecurityError, match="URL 为空"):
            validate_url_safety("")

    def test_none_url_rejected(self):
        """None URL被拒绝."""
        with pytest.raises(UrlSecurityError):
            validate_url_safety(None)

    def test_url_too_long_rejected(self):
        """超长URL被拒绝."""
        long_url = "https://example.com/" + "a" * 2048
        with pytest.raises(UrlSecurityError, match="URL 过长"):
            validate_url_safety(long_url)

    def test_invalid_scheme_rejected(self):
        """非法scheme被拒绝."""
        with pytest.raises(UrlSecurityError, match="不允许的 URL scheme"):
            validate_url_safety("ftp://example.com/file.mp3")

    def test_file_scheme_rejected(self):
        """file:// scheme被拒绝."""
        with pytest.raises(UrlSecurityError, match="不允许的 URL scheme"):
            validate_url_safety("file:///etc/passwd")

    def test_no_scheme_rejected(self):
        """无scheme被拒绝."""
        with pytest.raises(UrlSecurityError, match="不允许的 URL scheme"):
            validate_url_safety("example.com/file.mp3")

    def test_http_scheme_allowed(self):
        """http scheme允许（需要mock DNS）."""
        with patch("shared.url_security._check_ssrf_domain") as mock_check:
            result = validate_url_safety("http://example.com/audio.mp3")
            assert result == "http://example.com/audio.mp3"
            mock_check.assert_called_once()

    def test_https_scheme_allowed(self):
        """https scheme允许."""
        with patch("shared.url_security._check_ssrf_domain") as mock_check:
            result = validate_url_safety("https://example.com/audio.mp3")
            assert result == "https://example.com/audio.mp3"
            mock_check.assert_called_once()

    def test_scheme_case_insensitive(self):
        """scheme大小写不敏感."""
        with patch("shared.url_security._check_ssrf_domain") as mock_check:
            result = validate_url_safety("HTTPS://example.com/audio.mp3")
            assert result == "HTTPS://example.com/audio.mp3"

    def test_missing_hostname_rejected(self):
        """缺少主机名被拒绝."""
        with pytest.raises(UrlSecurityError, match="URL 缺少主机名"):
            validate_url_safety("https:///path/to/file")

    def test_non_standard_port_rejected(self):
        """非标准端口被拒绝."""
        with pytest.raises(UrlSecurityError, match="不允许的端口"):
            validate_url_safety("https://example.com:8080/audio.mp3")

    def test_port_22_rejected(self):
        """SSH端口被拒绝."""
        with pytest.raises(UrlSecurityError, match="不允许的端口"):
            validate_url_safety("https://example.com:22/")

    def test_port_80_allowed(self):
        """80端口允许."""
        with patch("shared.url_security._check_ssrf_domain") as mock_check:
            result = validate_url_safety("http://example.com:80/audio.mp3")
            assert result == "http://example.com:80/audio.mp3"

    def test_port_443_allowed(self):
        """443端口允许."""
        with patch("shared.url_security._check_ssrf_domain") as mock_check:
            result = validate_url_safety("https://example.com:443/audio.mp3")
            assert result == "https://example.com:443/audio.mp3"

    def test_default_http_port_implicit_80(self):
        """http默认端口隐含80，不触发端口校验."""
        with patch("shared.url_security._check_ssrf_domain") as mock_check:
            result = validate_url_safety("http://example.com/audio.mp3")
            assert "example.com" in result

    def test_default_https_port_implicit_443(self):
        """https默认端口隐含443，不触发端口校验."""
        with patch("shared.url_security._check_ssrf_domain") as mock_check:
            result = validate_url_safety("https://example.com/audio.mp3")
            assert "example.com" in result

    def test_purpose_parameter_logged(self):
        """purpose参数不影响校验结果."""
        with patch("shared.url_security._check_ssrf_domain"):
            result = validate_url_safety("https://example.com/a.mp3", purpose="tts_download")
            assert result == "https://example.com/a.mp3"

    def test_url_with_path_and_query(self):
        """带路径和query参数的URL正常通过."""
        with patch("shared.url_security._check_ssrf_domain"):
            url = "https://example.com/path/to/file.mp3?token=abc&expires=123"
            result = validate_url_safety(url)
            assert result == url


class TestInternalHostnames:
    """_check_internal_hostnames 内部主机名拦截."""

    def test_localhost_blocked(self):
        """localhost被拦截."""
        with pytest.raises(UrlSecurityError, match="禁止访问内部主机名"):
            _check_internal_hostnames("localhost")

    def test_localhost_case_insensitive(self):
        """大小写不敏感."""
        with pytest.raises(UrlSecurityError):
            _check_internal_hostnames("LOCALHOST")

    def test_localhost_localdomain_blocked(self):
        """localhost.localdomain被拦截."""
        with pytest.raises(UrlSecurityError):
            _check_internal_hostnames("localhost.localdomain")

    def test_metadata_blocked(self):
        """metadata被拦截."""
        with pytest.raises(UrlSecurityError):
            _check_internal_hostnames("metadata")

    def test_metadata_google_internal_blocked(self):
        """GCP元数据服务被拦截."""
        with pytest.raises(UrlSecurityError):
            _check_internal_hostnames("metadata.google.internal")

    def test_local_domain_suffix_blocked(self):
        """.local后缀被拦截."""
        with pytest.raises(UrlSecurityError, match="禁止访问内网域名"):
            _check_internal_hostnames("printer.local")

    def test_internal_domain_suffix_blocked(self):
        """.internal后缀被拦截."""
        with pytest.raises(UrlSecurityError, match="禁止访问内网域名"):
            _check_internal_hostnames("service.internal")

    def test_localdomain_suffix_blocked(self):
        """.localdomain后缀被拦截."""
        with pytest.raises(UrlSecurityError, match="禁止访问内网域名"):
            _check_internal_hostnames("host.localdomain")

    def test_public_domain_passes(self):
        """公网域名通过."""
        _check_internal_hostnames("example.com")  # 不抛异常即通过

    def test_subdomain_of_public_domain_passes(self):
        """公网域名子域名通过."""
        _check_internal_hostnames("cdn.example.com")  # 不抛异常即通过


class TestSSRFIPCheck:
    """_check_ssrf_ip SSRF IP检查."""

    def test_loopback_ipv4_blocked(self):
        """IPv4回环地址被拦截."""
        with pytest.raises(UrlSecurityError, match="回环地址"):
            _check_ssrf_ip(ipaddress.ip_address("127.0.0.1"))

    def test_loopback_ipv6_blocked(self):
        """IPv6回环地址被拦截."""
        with pytest.raises(UrlSecurityError, match="回环地址"):
            _check_ssrf_ip(ipaddress.ip_address("::1"))

    def test_private_ip_10_blocked(self):
        """10.x.x.x私有地址被拦截."""
        with pytest.raises(UrlSecurityError, match="内网地址"):
            _check_ssrf_ip(ipaddress.ip_address("10.0.0.1"))

    def test_private_ip_172_blocked(self):
        """172.16.x.x私有地址被拦截."""
        with pytest.raises(UrlSecurityError, match="内网地址"):
            _check_ssrf_ip(ipaddress.ip_address("172.16.0.1"))

    def test_private_ip_192_blocked(self):
        """192.168.x.x私有地址被拦截."""
        with pytest.raises(UrlSecurityError, match="内网地址"):
            _check_ssrf_ip(ipaddress.ip_address("192.168.1.1"))

    def test_link_local_blocked(self):
        """链路本地地址被拦截（is_private 先命中也可以，只要被拦就行）."""
        with pytest.raises(UrlSecurityError):
            _check_ssrf_ip(ipaddress.ip_address("169.254.169.254"))

    def test_multicast_blocked(self):
        """组播地址被拦截."""
        with pytest.raises(UrlSecurityError, match="组播地址"):
            _check_ssrf_ip(ipaddress.ip_address("224.0.0.1"))

    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_ip_blocked(self):
        """保留地址被拦截（240.0.0.0/4 属于保留段）."""
        with pytest.raises(UrlSecurityError):
            _check_ssrf_ip(ipaddress.ip_address("240.0.0.1"))

    def test_public_ip_passes(self):
        """公网IP通过（当ALLOW_DIRECT_IP=true时）."""
        # 注意：公网IP在_check_ssrf_ip层面不拦截，拦截在validate_url_safety的ALLOW_DIRECT_IP层
        _check_ssrf_ip(ipaddress.ip_address("8.8.8.8"))  # 不抛异常即通过

    def test_public_ip_another_passes(self):
        """另一个公网IP通过."""
        _check_ssrf_ip(ipaddress.ip_address("1.1.1.1"))  # 不抛异常即通过

    def test_ipv6_public_passes(self):
        """IPv6公网地址通过."""
        _check_ssrf_ip(ipaddress.ip_address("2001:4860:4860::8888"))  # 不抛异常即通过


class TestDirectIPBlocking:
    """直接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_ipv6_blocked_by_default(self):
        """默认禁止直接IPv6访问."""
        with pytest.raises(UrlSecurityError, match="禁止直接 IP 访问"):
            validate_url_safety("https://[2001:4860:4860::8888]/audio.mp3")

    def test_direct_ip_allowed_when_flag_set(self):
        """ALLOW_DIRECT_IP=true时允许公网IP直接访问."""
        with patch("shared.url_security.ALLOW_DIRECT_IP", True):
            # 公网IP应该通过
            with patch("shared.url_security._check_ssrf_ip") as mock_check:
                result = validate_url_safety("https://8.8.8.8/audio.mp3")
                assert result == "https://8.8.8.8/audio.mp3"
                mock_check.assert_called_once()

    def test_direct_private_ip_blocked_even_with_flag(self):
        """即使ALLOW_DIRECT_IP=true，私有IP仍被SSRF检查拦截."""
        with patch("shared.url_security.ALLOW_DIRECT_IP", True):
            with pytest.raises(UrlSecurityError, match="内网地址"):
                validate_url_safety("https://192.168.1.1/audio.mp3")


class TestTrustedDomain:
    """_is_trusted_domain 可信域名匹配."""

    def test_exact_match(self):
        """精确匹配."""
        with patch("shared.url_security.TRUSTED_DOMAINS", {"example.com"}):
            assert _is_trusted_domain("example.com") is True

    def test_subdomain_match(self):
        """子域名匹配."""
        with patch("shared.url_security.TRUSTED_DOMAINS", {"example.com"}):
            assert _is_trusted_domain("cdn.example.com") is True

    def test_nested_subdomain_match(self):
        """多级子域名匹配."""
        with patch("shared.url_security.TRUSTED_DOMAINS", {"example.com"}):
            assert _is_trusted_domain("a.b.c.example.com") is True

    def test_no_match(self):
        """不匹配的域名."""
        with patch("shared.url_security.TRUSTED_DOMAINS", {"example.com"}):
            assert _is_trusted_domain("other.com") is False

    def test_domain_containing_but_not_subdomain(self):
        """域名包含但不是子域名."""
        with patch("shared.url_security.TRUSTED_DOMAINS", {"example.com"}):
            assert _is_trusted_domain("fake-example.com") is False

    def test_case_insensitive(self):
        """大小写不敏感."""
        with patch("shared.url_security.TRUSTED_DOMAINS", {"Example.COM"}):
            assert _is_trusted_domain("cdn.example.com") is True

    def test_empty_trusted_domains(self):
        """空白名单时全部返回False."""
        with patch("shared.url_security.TRUSTED_DOMAINS", set()):
            assert _is_trusted_domain("anything.com") is False

    def test_multiple_trusted_domains(self):
        """多个可信域名."""
        with patch("shared.url_security.TRUSTED_DOMAINS", {"example.com", "aliyuncs.com"}):
            assert _is_trusted_domain("bucket.oss-cn-hangzhou.aliyuncs.com") is True
            assert _is_trusted_domain("cdn.example.com") is True
            assert _is_trusted_domain("unknown.com") is False


class TestIsUrlSafe:
    """is_url_safe 便捷函数."""

    def test_safe_url_returns_true(self):
        """安全URL返回True."""
        with patch("shared.url_security._check_ssrf_domain"):
            assert is_url_safe("https://example.com/a.mp3") is True

    def test_unsafe_url_returns_false(self):
        """不安全URL返回False."""
        assert is_url_safe("ftp://example.com/a.mp3") is False

    def test_empty_url_returns_false(self):
        """空URL返回False."""
        assert is_url_safe("") is False

    def test_does_not_raise(self):
        """不抛异常."""
        try:
            is_url_safe("not a url at all")
        except Exception:
            pytest.fail("is_url_safe should not raise exceptions")


class TestMagicNumberValidation:
    """_validate_magic_number 文件魔数校验."""

    def test_valid_mp3_id3v2(self):
        """MP3 ID3v2标签魔数通过."""
        with tempfile.NamedTemporaryFile(suffix=".mp3", delete=False) as f:
            f.write(b"ID3" + b"\x00" * 100)
            tmp_path = f.name
        try:
            _validate_magic_number(tmp_path, {"audio/mpeg"})  # 不抛异常即通过
        finally:
            os.unlink(tmp_path)

    def test_valid_wav(self):
        """WAV文件魔数通过."""
        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
            # RIFF + WAVE
            f.write(b"RIFF" + b"\x00" * 4 + b"WAVE" + b"\x00" * 100)
            tmp_path = f.name
        try:
            _validate_magic_number(tmp_path, {"audio/wav"})  # 不抛异常即通过
        finally:
            os.unlink(tmp_path)

    def test_valid_png(self):
        """PNG图片魔数通过."""
        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as f:
            f.write(b"\x89PNG\r\n\x1a\n" + b"\x00" * 100)
            tmp_path = f.name
        try:
            _validate_magic_number(tmp_path, {"image/png"})  # 不抛异常即通过
        finally:
            os.unlink(tmp_path)

    def test_valid_jpeg(self):
        """JPEG图片魔数通过."""
        with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as f:
            f.write(b"\xff\xd8\xff" + b"\x00" * 100)
            tmp_path = f.name
        try:
            _validate_magic_number(tmp_path, {"image/jpeg"})  # 不抛异常即通过
        finally:
            os.unlink(tmp_path)

    def test_valid_gif87a(self):
        """GIF87a通过."""
        with tempfile.NamedTemporaryFile(suffix=".gif", delete=False) as f:
            f.write(b"GIF87a" + b"\x00" * 100)
            tmp_path = f.name
        try:
            _validate_magic_number(tmp_path, {"image/gif"})  # 不抛异常即通过
        finally:
            os.unlink(tmp_path)

    def test_valid_gif89a(self):
        """GIF89a通过."""
        with tempfile.NamedTemporaryFile(suffix=".gif", delete=False) as f:
            f.write(b"GIF89a" + b"\x00" * 100)
            tmp_path = f.name
        try:
            _validate_magic_number(tmp_path, {"image/gif"})  # 不抛异常即通过
        finally:
            os.unlink(tmp_path)

    def test_valid_ogg(self):
        """OGG音频通过."""
        with tempfile.NamedTemporaryFile(suffix=".ogg", delete=False) as f:
            f.write(b"OggS" + b"\x00" * 100)
            tmp_path = f.name
        try:
            _validate_magic_number(tmp_path, {"audio/ogg"})  # 不抛异常即通过
        finally:
            os.unlink(tmp_path)

    def test_invalid_magic_number_rejected(self):
        """错误魔数被拒绝."""
        with tempfile.NamedTemporaryFile(suffix=".bin", delete=False) as f:
            f.write(b"NOT_A_VALID_FILE" + b"\x00" * 100)
            tmp_path = f.name
        try:
            with pytest.raises(UrlSecurityError, match="文件魔数与允许的 MIME 类型不匹配"):
                _validate_magic_number(tmp_path, {"audio/mpeg", "image/png"})
        finally:
            os.unlink(tmp_path)

    def test_empty_file_rejected(self):
        """空文件被拒绝."""
        with tempfile.NamedTemporaryFile(suffix=".bin", delete=False) as f:
            tmp_path = f.name
        try:
            with pytest.raises(UrlSecurityError, match="文件为空"):
                _validate_magic_number(tmp_path, {"audio/mpeg"})
        finally:
            os.unlink(tmp_path)

    def test_unknown_mime_skipped(self):
        """未知MIME类型跳过校验."""
        with tempfile.NamedTemporaryFile(suffix=".bin", delete=False) as f:
            f.write(b"random data here")
            tmp_path = f.name
        try:
            # 未知MIME类型没有对应的魔数，应该跳过不阻断
            _validate_magic_number(tmp_path, {"application/x-unknown-type"})  # 不抛异常即通过
        finally:
            os.unlink(tmp_path)

    def test_empty_allowed_mime_skipped(self):
        """空allowed_mime_types集合直接通过（无魔数可比对）."""
        with tempfile.NamedTemporaryFile(suffix=".bin", delete=False) as f:
            f.write(b"anything")
            tmp_path = f.name
        try:
            _validate_magic_number(tmp_path, set())  # 不抛异常即通过
        finally:
            os.unlink(tmp_path)

    def test_multiple_allowed_types_one_matches(self):
        """多个允许类型，只要一个匹配就通过."""
        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as f:
            f.write(b"\x89PNG\r\n\x1a\n" + b"\x00" * 100)
            tmp_path = f.name
        try:
            _validate_magic_number(tmp_path, {"audio/mpeg", "image/png", "image/jpeg"})
        finally:
            os.unlink(tmp_path)

    def test_mp4_ftyp_magic(self):
        """MP4 ftyp魔数通过."""
        with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as f:
            # ftyp在偏移4的位置
            f.write(b"\x00\x00\x00\x20" + b"ftyp" + b"mp42" + b"\x00" * 100)
            tmp_path = f.name
        try:
            _validate_magic_number(tmp_path, {"video/mp4"})  # 不抛异常即通过
        finally:
            os.unlink(tmp_path)

    def test_webp_magic(self):
        """WebP魔数通过."""
        with tempfile.NamedTemporaryFile(suffix=".webp", delete=False) as f:
            f.write(b"RIFF" + b"\x00" * 4 + b"WEBP" + b"\x00" * 100)
            tmp_path = f.name
        try:
            _validate_magic_number(tmp_path, {"image/webp"})  # 不抛异常即通过
        finally:
            os.unlink(tmp_path)

    def test_flac_magic(self):
        """FLAC魔数通过."""
        with tempfile.NamedTemporaryFile(suffix=".flac", delete=False) as f:
            f.write(b"fLaC" + b"\x00" * 100)
            tmp_path = f.name
        try:
            _validate_magic_number(tmp_path, {"audio/flac"})  # 不抛异常即通过
        finally:
            os.unlink(tmp_path)

    def test_file_too_short_for_signature(self):
        """文件太短，不够签名长度时不匹配."""
        with tempfile.NamedTemporaryFile(suffix=".bin", delete=False) as f:
            f.write(b"ID")  # 只有2字节，不够ID3的3字节
            tmp_path = f.name
        try:
            with pytest.raises(UrlSecurityError, match="文件魔数"):
                _validate_magic_number(tmp_path, {"audio/mpeg"})
        finally:
            os.unlink(tmp_path)