Livrare LOT 1 - Didi

ai_platform/modules/web/tests/test_validation.py

@@ -0,0 +1,146 @@
+"""Tests for URL validation (SSRF protection)."""
+
+from unittest.mock import AsyncMock, patch
+
+import pytest
+
+from web.validation import validate_url, validate_url_dns, validate_urls_async
+
+
+class TestValidateUrl:
+    """Tests for validate_url."""
+
+    def test_valid_https_url(self) -> None:
+        """Test that a valid HTTPS URL passes."""
+        result = validate_url("https://example.com/article")
+        assert result == "https://example.com/article"
+
+    def test_valid_http_url(self) -> None:
+        """Test that a valid HTTP URL passes."""
+        result = validate_url("http://example.com/page")
+        assert result == "http://example.com/page"
+
+    def test_rejects_ftp_scheme(self) -> None:
+        """Test that FTP scheme is rejected."""
+        with pytest.raises(ValueError, match="not allowed"):
+            validate_url("ftp://example.com/file")
+
+    def test_rejects_file_scheme(self) -> None:
+        """Test that file:// scheme is rejected."""
+        with pytest.raises(ValueError, match="not allowed"):
+            validate_url("file:///etc/passwd")
+
+    def test_rejects_javascript_scheme(self) -> None:
+        """Test that javascript: scheme is rejected."""
+        with pytest.raises(ValueError, match="not allowed"):
+            validate_url("javascript:alert(1)")
+
+    def test_rejects_localhost(self) -> None:
+        """Test that localhost is rejected."""
+        with pytest.raises(ValueError, match="not allowed"):
+            validate_url("http://localhost/secret")
+
+    def test_rejects_127_0_0_1(self) -> None:
+        """Test that 127.0.0.1 is rejected."""
+        with pytest.raises(ValueError, match="private/reserved"):
+            validate_url("http://127.0.0.1/secret")
+
+    def test_rejects_private_ip_10(self) -> None:
+        """Test that 10.x.x.x private IPs are rejected."""
+        with pytest.raises(ValueError, match="private/reserved"):
+            validate_url("http://10.0.0.1/internal")
+
+    def test_rejects_private_ip_192_168(self) -> None:
+        """Test that 192.168.x.x private IPs are rejected."""
+        with pytest.raises(ValueError, match="private/reserved"):
+            validate_url("http://192.168.1.1/admin")
+
+    def test_rejects_private_ip_172_16(self) -> None:
+        """Test that 172.16.x.x private IPs are rejected."""
+        with pytest.raises(ValueError, match="private/reserved"):
+            validate_url("http://172.16.0.1/internal")
+
+    def test_rejects_metadata_endpoint(self) -> None:
+        """Test that cloud metadata endpoint is rejected."""
+        with pytest.raises(ValueError, match="not allowed"):
+            validate_url("http://169.254.169.254/latest/meta-data")
+
+    def test_rejects_no_hostname(self) -> None:
+        """Test that URL without hostname is rejected."""
+        with pytest.raises(ValueError, match="hostname"):
+            validate_url("http:///path")
+
+    def test_rejects_empty_scheme(self) -> None:
+        """Test that URL without scheme is rejected."""
+        with pytest.raises(ValueError, match="not allowed"):
+            validate_url("example.com/page")
+
+
+class TestValidateUrlDns:
+    """Tests for async DNS-based SSRF validation."""
+
+    @pytest.mark.asyncio
+    async def test_public_url_passes(self) -> None:
+        """Test that a URL resolving to a public IP passes."""
+        # Mock getaddrinfo to return a public IP
+        mock_result = [(2, 1, 6, "", ("93.184.216.34", 0))]
+        with patch("web.validation.asyncio.get_running_loop") as mock_loop:
+            mock_loop.return_value.getaddrinfo = AsyncMock(return_value=mock_result)
+            result = await validate_url_dns("https://example.com/article")
+        assert result == "https://example.com/article"
+
+    @pytest.mark.asyncio
+    async def test_rejects_dns_to_private_ip(self) -> None:
+        """Test that a hostname resolving to a private IP is rejected."""
+        mock_result = [(2, 1, 6, "", ("127.0.0.1", 0))]
+        with patch("web.validation.asyncio.get_running_loop") as mock_loop:
+            mock_loop.return_value.getaddrinfo = AsyncMock(return_value=mock_result)
+            with pytest.raises(ValueError, match="private/reserved"):
+                await validate_url_dns("https://evil.example.com/steal")
+
+    @pytest.mark.asyncio
+    async def test_dns_failure_rejects_url(self) -> None:
+        """Test that DNS resolution failure blocks the URL (fail closed)."""
+        import socket
+
+        with patch("web.validation.asyncio.get_running_loop") as mock_loop:
+            mock_loop.return_value.getaddrinfo = AsyncMock(
+                side_effect=socket.gaierror("Name or service not known")
+            )
+            with pytest.raises(ValueError, match="could not be resolved"):
+                await validate_url_dns("https://nonexistent.example.com/page")
+
+    @pytest.mark.asyncio
+    async def test_no_hostname_passes(self) -> None:
+        """Test that URL with no hostname returns immediately."""
+        # This shouldn't happen in practice (validate_url catches it)
+        # but test the guard clause
+        result = await validate_url_dns("http:///path")
+        assert result == "http:///path"
+
+
+class TestValidateUrlsAsync:
+    """Tests for batch async URL validation."""
+
+    @pytest.mark.asyncio
+    async def test_validates_multiple_urls(self) -> None:
+        """Test that multiple URLs are validated in parallel."""
+        mock_result = [(2, 1, 6, "", ("93.184.216.34", 0))]
+        with patch("web.validation.asyncio.get_running_loop") as mock_loop:
+            mock_loop.return_value.getaddrinfo = AsyncMock(return_value=mock_result)
+            urls = [
+                "https://example.com/a",
+                "https://example.com/b",
+                "https://example.com/c",
+            ]
+            result = await validate_urls_async(urls)
+        assert result == urls
+
+    @pytest.mark.asyncio
+    async def test_raises_on_private_ip_in_batch(self) -> None:
+        """Test that a private IP in a batch raises."""
+        mock_result = [(2, 1, 6, "", ("10.0.0.1", 0))]
+        with patch("web.validation.asyncio.get_running_loop") as mock_loop:
+            mock_loop.return_value.getaddrinfo = AsyncMock(return_value=mock_result)
+            with pytest.raises(ValueError, match="private/reserved"):
+                await validate_urls_async(["https://evil.example.com/steal"])