archived-MoviePilot/tests/test_security_utils.py

import socket
from unittest import TestCase
from unittest.mock import patch

from app.utils.security import SecurityUtils


class SecurityUtilsTest(TestCase):

    def test_is_safe_url_keeps_default_allowlist_behavior(self):
        """
        默认 URL 校验保持历史 allowlist 行为，避免影响非代理调用方。
        """
        self.assertTrue(
            SecurityUtils.is_safe_url(
                "http://192.168.1.50:8096/secret.png",
                {"http://192.168.1.50:8096"},
            )
        )

    def test_is_safe_url_blocks_private_literal_ip_when_enabled(self):
        """
        启用 SSRF 防护时，即使内网 IP 命中 allowlist 也不能放行。
        """
        self.assertFalse(
            SecurityUtils.is_safe_url(
                "http://192.168.1.50:8096/secret.png",
                {"http://192.168.1.50:8096"},
                block_private=True,
            )
        )

    def test_is_safe_url_blocks_loopback_dns_result_when_enabled(self):
        """
        主机名解析到回环地址时必须拒绝，防止通过域名绕过内网地址拦截。
        """
        with patch(
            "app.utils.security.socket.getaddrinfo",
            return_value=[
                (
                    socket.AF_INET,
                    socket.SOCK_STREAM,
                    0,
                    "",
                    ("127.0.0.1", 0),
                )
            ],
        ):
            self.assertFalse(
                SecurityUtils.is_safe_url(
                    "http://internal.example.com/secret.png",
                    {"example.com"},
                    block_private=True,
                )
            )

    def test_is_safe_url_blocks_mixed_public_and_private_dns_results(self):
        """
        同一域名只要存在任一非公网解析结果，就不能作为图片代理目标。
        """
        with patch(
            "app.utils.security.socket.getaddrinfo",
            return_value=[
                (
                    socket.AF_INET,
                    socket.SOCK_STREAM,
                    0,
                    "",
                    ("93.184.216.34", 0),
                ),
                (
                    socket.AF_INET,
                    socket.SOCK_STREAM,
                    0,
                    "",
                    ("10.0.0.8", 0),
                ),
            ],
        ):
            self.assertFalse(
                SecurityUtils.is_safe_url(
                    "https://assets.example.com/poster.jpg",
                    {"example.com"},
                    block_private=True,
                )
            )

    def test_is_safe_url_allows_public_dns_result_when_enabled(self):
        """
        域名解析结果全部为公网地址且命中 allowlist 时继续允许访问。
        """
        with patch(
            "app.utils.security.socket.getaddrinfo",
            return_value=[
                (
                    socket.AF_INET,
                    socket.SOCK_STREAM,
                    0,
                    "",
                    ("93.184.216.34", 0),
                )
            ],
        ):
            self.assertTrue(
                SecurityUtils.is_safe_url(
                    "https://assets.example.com/poster.jpg",
                    {"example.com"},
                    block_private=True,
                )
            )

    def test_is_safe_url_rejects_dns_resolution_failure_when_enabled(self):
        """
        SSRF 防护无法确认目标地址时按失败处理，避免解析异常时继续请求。
        """
        with patch(
            "app.utils.security.socket.getaddrinfo",
            side_effect=socket.gaierror,
        ):
            self.assertFalse(
                SecurityUtils.is_safe_url(
                    "https://assets.example.com/poster.jpg",
                    {"example.com"},
                    block_private=True,
                )
            )