xiaoxia-saas/tests/unit/test_url_security.py

"""URL 安全校验工具单元测试 — SSRF 防护."""

from __future__ import annotations

import os
import sys
import unittest

sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", "apps", "worker"))

from video_processing.url_security import (  # noqa: E402
    UrlSecurityError,
    is_url_safe,
    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"))


if __name__ == "__main__":
    unittest.main()