xiaoxia-saas/tests/unit/test_url_security_domain.py

"""URL 安全校验纯逻辑单元测试 — wave128."""

import pytest

from packages.domain.url_security import (
    ALLOWED_PORTS,
    ALLOWED_SCHEMES,
    MAGIC_NUMBERS,
    MAX_URL_LENGTH,
    UrlSecurityError,
    check_internal_hostname,
    check_ssrf_ip,
    is_ip_address,
    is_trusted_domain,
    is_url_basic_safe,
    validate_magic_number,
    validate_url_basic,
)

# ── 常量校验 ────────────────────────────────────────────────────────────────


class TestConstants:
    def test_allowed_schemes(self):
        assert "http" in ALLOWED_SCHEMES
        assert "https" in ALLOWED_SCHEMES

    def test_allowed_ports(self):
        assert 80 in ALLOWED_PORTS
        assert 443 in ALLOWED_PORTS

    def test_max_url_length(self):
        assert MAX_URL_LENGTH == 2048

    def test_magic_numbers_has_common_formats(self):
        assert "audio/mpeg" in MAGIC_NUMBERS
        assert "image/png" in MAGIC_NUMBERS
        assert "video/mp4" in MAGIC_NUMBERS


# ── 内部主机名检查 ──────────────────────────────────────────────────────────


class TestCheckInternalHostname:
    @pytest.mark.parametrize(
        "hostname",
        [
            "localhost",
            "LOCALHOST",
            "LocalHost",
            "localhost.localdomain",
            "ip6-localhost",
            "ip6-loopback",
            "metadata",
            "metadata.google.internal",
            "169.254.169.254",
        ],
    )
    def test_internal_hostnames_rejected(self, hostname):
        with pytest.raises(UrlSecurityError, match="禁止访问内部主机名"):
            check_internal_hostname(hostname)

    @pytest.mark.parametrize(
        "hostname",
        [
            "foo.local",
            "bar.internal",
            "baz.localdomain",
            "sub.foo.local",
        ],
    )
    def test_internal_domain_suffixes_rejected(self, hostname):
        with pytest.raises(UrlSecurityError, match="禁止访问内网域名"):
            check_internal_hostname(hostname)

    @pytest.mark.parametrize(
        "hostname",
        [
            "example.com",
            "www.google.com",
            "oss-cn-hangzhou.aliyuncs.com",
            "123.45.67.89",
        ],
    )
    def test_normal_hostnames_allowed(self, hostname):
        check_internal_hostname("example.com")  # 不抛异常即通过


# ── 可信域名匹配 ────────────────────────────────────────────────────────────


class TestIsTrustedDomain:
    def test_empty_trusted_always_true(self):
        assert is_trusted_domain("anything.com", set()) is True

    def test_exact_match(self):
        trusted = {"example.com", "foo.bar"}
        assert is_trusted_domain("example.com", trusted) is True
        assert is_trusted_domain("foo.bar", trusted) is True

    def test_exact_no_match(self):
        trusted = {"example.com"}
        assert is_trusted_domain("other.com", trusted) is False

    def test_subdomain_match(self):
        trusted = {"example.com"}
        assert is_trusted_domain("sub.example.com", trusted) is True
        assert is_trusted_domain("a.b.example.com", trusted) is True

    def test_subdomain_partial_no_match(self):
        trusted = {"example.com"}
        # fakeexample.com 不是 example.com 的子域名
        assert is_trusted_domain("fakeexample.com", trusted) is False

    def test_case_insensitive(self):
        trusted = {"Example.COM"}
        assert is_trusted_domain("example.com", trusted) is True
        assert is_trusted_domain("SUB.Example.COM", trusted) is True


# ── IP SSRF 检查 ────────────────────────────────────────────────────────────


class TestCheckSrfIp:
    @pytest.mark.parametrize("ip", ["127.0.0.1", "127.1.2.3", "::1"])
    def test_loopback_rejected(self, ip):
        with pytest.raises(UrlSecurityError, match="回环"):
            check_ssrf_ip(ip)

    @pytest.mark.parametrize(
        "ip",
        [
            "10.0.0.1",
            "10.255.255.255",
            "172.16.0.1",
            "172.31.255.255",
            "192.168.1.1",
            "192.168.0.1",
            "fd00::1",  # IPv6 unique local
        ],
    )
    def test_private_rejected(self, ip):
        with pytest.raises(UrlSecurityError, match="内网"):
            check_ssrf_ip(ip)

    @pytest.mark.parametrize("ip", ["169.254.1.1", "169.254.169.254", "fe80::1"])
    def test_link_local_rejected(self, ip):
        with pytest.raises(UrlSecurityError, match="链路本地"):
            check_ssrf_ip(ip)

    @pytest.mark.parametrize("ip", ["224.0.0.1", "239.255.255.255", "ff00::1"])
    def test_multicast_rejected(self, ip):
        with pytest.raises(UrlSecurityError, match="组播"):
            check_ssrf_ip(ip)

    @pytest.mark.parametrize("ip", ["0.0.0.0", "::"])
    def test_unspecified_rejected(self, ip):
        with pytest.raises(UrlSecurityError, match="未指定"):
            check_ssrf_ip(ip)

    def test_reserved_rejected(self):
        with pytest.raises(UrlSecurityError):
            check_ssrf_ip("240.0.0.1")  # 保留地址段

    @pytest.mark.parametrize(
        "ip",
        [
            "8.8.8.8",
            "1.1.1.1",
            "223.5.5.5",
            "2001:4860:4860::8888",
        ],
    )
    def test_public_ip_allowed(self, ip):
        check_ssrf_ip(ip)  # 不抛异常即通过

    def test_invalid_ip_raises_value_error(self):
        with pytest.raises(ValueError):
            check_ssrf_ip("not-an-ip")


# ── IP 地址判断 ─────────────────────────────────────────────────────────────


class TestIsIpAddress:
    @pytest.mark.parametrize(
        "host",
        [
            "127.0.0.1",
            "8.8.8.8",
            "192.168.1.1",
            "::1",
            "2001:db8::1",
            "fe80::1",
        ],
    )
    def test_ip_addresses(self, host):
        assert is_ip_address(host) is True

    @pytest.mark.parametrize(
        "host",
        [
            "example.com",
            "www.google.com",
            "localhost",
            "not-an-ip",
            "",
        ],
    )
    def test_not_ip_addresses(self, host):
        assert is_ip_address(host) is False


# ── URL 基础校验 ────────────────────────────────────────────────────────────


class TestValidateUrlBasic:
    def test_normal_http_url_passes(self):
        result = validate_url_basic("http://example.com/file.txt")
        assert result == "http://example.com/file.txt"

    def test_normal_https_url_passes(self):
        result = validate_url_basic("https://www.example.com/path?q=1")
        assert result == "https://www.example.com/path?q=1"

    def test_standard_port_80_passes(self):
        validate_url_basic("http://example.com:80/file")

    def test_standard_port_443_passes(self):
        validate_url_basic("https://example.com:443/file")

    def test_empty_url_rejected(self):
        with pytest.raises(UrlSecurityError, match="URL 为空"):
            validate_url_basic("")

    def test_none_url_rejected(self):
        with pytest.raises(UrlSecurityError, match="URL 为空"):
            validate_url_basic(None)  # type: ignore

    def test_too_long_url_rejected(self):
        long_url = "https://example.com/" + "a" * 2100
        with pytest.raises(UrlSecurityError, match="URL 过长"):
            validate_url_basic(long_url)

    @pytest.mark.parametrize(
        "url",
        [
            "ftp://example.com/file",
            "file:///etc/passwd",
            "javascript:alert(1)",
            "data:text/html,<h1>hi</h1>",
        ],
    )
    def test_bad_scheme_rejected(self, url):
        with pytest.raises(UrlSecurityError, match="不允许的 URL scheme"):
            validate_url_basic(url)

    def test_missing_hostname_rejected(self):
        with pytest.raises(UrlSecurityError, match="URL 缺少主机名"):
            validate_url_basic("http:///path")

    @pytest.mark.parametrize(
        "url",
        [
            "http://localhost/test",
            "http://metadata/test",
            "http://foo.local/test",
        ],
    )
    def test_internal_hostname_rejected(self, url):
        with pytest.raises(UrlSecurityError):
            validate_url_basic(url)

    @pytest.mark.parametrize(
        "url",
        [
            "http://example.com:8080/file",
            "http://example.com:22/file",
            "http://example.com:3306/file",
        ],
    )
    def test_non_standard_port_rejected(self, url):
        with pytest.raises(UrlSecurityError, match="不允许的端口"):
            validate_url_basic(url)

    def test_direct_ip_rejected_by_default(self):
        with pytest.raises(UrlSecurityError, match="禁止直接 IP 访问"):
            validate_url_basic("http://8.8.8.8/file")

    def test_direct_ip_allowed_with_flag_public(self):
        result = validate_url_basic("http://8.8.8.8/file", allow_direct_ip=True)
        assert result == "http://8.8.8.8/file"

    def test_direct_ip_allowed_flag_but_private_still_rejected(self):
        with pytest.raises(UrlSecurityError, match="内网地址"):
            validate_url_basic("http://10.0.0.1/file", allow_direct_ip=True)

    def test_direct_ip_loopback_rejected(self):
        with pytest.raises(UrlSecurityError):
            validate_url_basic("http://127.0.0.1/test", allow_direct_ip=True)

    def test_trusted_domains_whitelist_pass(self):
        trusted = {"example.com"}
        result = validate_url_basic("https://example.com/file", trusted_domains=trusted)
        assert result == "https://example.com/file"

    def test_trusted_domains_subdomain_pass(self):
        trusted = {"example.com"}
        result = validate_url_basic("https://cdn.example.com/file", trusted_domains=trusted)
        assert result == "https://cdn.example.com/file"

    def test_trusted_domains_not_in_list_rejected(self):
        trusted = {"example.com"}
        with pytest.raises(UrlSecurityError, match="不在可信白名单"):
            validate_url_basic("https://other.com/file", trusted_domains=trusted)

    def test_case_insensitive_scheme(self):
        # 大写 HTTP 也应该通过（我们用 .lower() 检查）
        result = validate_url_basic("HTTP://example.com/file")
        assert "HTTP://example.com/file" == result


# ── is_url_basic_safe 便捷函数 ──────────────────────────────────────────────


class TestIsUrlBasicSafe:
    def test_safe_url_returns_true(self):
        assert is_url_basic_safe("https://example.com/file") is True

    def test_unsafe_url_returns_false(self):
        assert is_url_basic_safe("http://localhost/test") is False

    def test_empty_returns_false(self):
        assert is_url_basic_safe("") is False

    def test_with_trusted_domains(self):
        trusted = {"allowed.com"}
        assert is_url_basic_safe("https://allowed.com/x", trusted_domains=trusted) is True
        assert is_url_basic_safe("https://other.com/x", trusted_domains=trusted) is False


# ── 魔数校验 ────────────────────────────────────────────────────────────────


class TestValidateMagicNumber:
    def test_mp3_id3_header(self):
        data = b"ID3" + b"\x00" * 100
        validate_magic_number(data, {"audio/mpeg"})  # 不抛异常

    def test_mp3_frame_sync(self):
        data = b"\xff\xfb" + b"\x00" * 100
        validate_magic_number(data, {"audio/mpeg"})

    def test_wav_header(self):
        data = b"RIFF" + b"\x00" * 4 + b"WAVE" + b"\x00" * 100
        validate_magic_number(data, {"audio/wav"})

    def test_png_header(self):
        data = b"\x89PNG\r\n\x1a\n" + b"\x00" * 100
        validate_magic_number(data, {"image/png"})

    def test_jpeg_header(self):
        data = b"\xff\xd8\xff" + b"\x00" * 100
        validate_magic_number(data, {"image/jpeg"})

    def test_gif87a_header(self):
        data = b"GIF87a" + b"\x00" * 100
        validate_magic_number(data, {"image/gif"})

    def test_gif89a_header(self):
        data = b"GIF89a" + b"\x00" * 100
        validate_magic_number(data, {"image/gif"})

    def test_mp4_ftyp_header(self):
        data = b"\x00\x00\x00\x20ftypisom" + b"\x00" * 100
        validate_magic_number(data, {"video/mp4"})

    def test_ogg_header(self):
        data = b"OggS" + b"\x00" * 100
        validate_magic_number(data, {"audio/ogg"})

    def test_flac_header(self):
        data = b"fLaC" + b"\x00" * 100
        validate_magic_number(data, {"audio/flac"})

    def test_webp_header(self):
        data = b"RIFF" + b"\x00" * 4 + b"WEBP" + b"\x00" * 100
        validate_magic_number(data, {"image/webp"})

    def test_bmp_header(self):
        data = b"BM" + b"\x00" * 100
        validate_magic_number(data, {"image/bmp"})

    def test_empty_file_rejected(self):
        with pytest.raises(UrlSecurityError, match="文件为空"):
            validate_magic_number(b"", {"image/png"})

    def test_mismatched_magic_rejected(self):
        data = b"NOTAPNG" + b"\x00" * 100
        with pytest.raises(UrlSecurityError, match="魔数与允许的 MIME 类型不匹配"):
            validate_magic_number(data, {"image/png"})

    def test_multiple_allowed_types_one_matches(self):
        data = b"\x89PNG\r\n\x1a\n" + b"\x00" * 100
        # 多个允许类型，只要有一个匹配就通过
        validate_magic_number(data, {"image/png", "image/jpeg", "image/gif"})

    def test_no_known_magic_skips_validation(self):
        # 自定义 MIME 类型没有已知魔数，跳过校验不阻断
        validate_magic_number(b"random data", {"application/x-custom"})

    def test_too_short_header_no_match(self):
        # 文件头太短，无法匹配需要 8 字节偏移的格式
        data = b"RIFF"  # 只有 4 字节，不够 offset 8 的 WAVE 匹配
        with pytest.raises(UrlSecurityError, match="魔数"):
            validate_magic_number(data, {"audio/wav"})

    def test_error_message_contains_mime_and_header(self):
        with pytest.raises(UrlSecurityError) as exc_info:
            validate_magic_number(b"XXXXYYY", {"image/png"})
        msg = str(exc_info.value)
        assert "image/png" in msg
        assert "文件头前16字节" in msg