"""Test DNS rebinding protection in API Request and URL components.

This test suite verifies that the DNS pinning implementation prevents
DNS rebinding attacks that could bypass SSRF protection in both the
API Request component and the URL component.
"""
# ruff: noqa: ARG001, SIM117

import os
import socket
from unittest.mock import patch

import httpcore
import httpx
import pytest
from lfx.components.data_source.api_request import APIRequestComponent
from lfx.components.data_source.url import URLComponent
from lfx.schema import Data


class TestDNSRebindingProtection:
    """Test DNS rebinding attack prevention through DNS pinning."""

    @pytest.fixture
    def component(self):
        """Create a basic API request component."""
        return APIRequestComponent(
            url_input="http://rebinding.test:8080/api",
            method="GET",
            headers=[],
            body=[],
            timeout=30,
            follow_redirects=False,
            save_to_file=False,
            include_httpx_metadata=False,
            mode="URL",
            curl_input="",
            query_params={},
        )

    @pytest.mark.asyncio
    async def test_dns_pinning_prevents_rebinding_attack(self, component):
        """Test that DNS pinning prevents DNS rebinding attacks.

        This test simulates a DNS rebinding attack where:
        1. First DNS lookup (validation): returns public IP (8.8.8.8)
        2. Second DNS lookup (httpx): would return localhost (127.0.0.1)

        With DNS pinning, the second lookup should NOT happen - the validated
        IP from the first lookup should be used directly at the network layer.
        """
        call_count = 0
        connected_to_ip = None

        def mock_getaddrinfo(_hostname, _port, *_args, **_kwargs):
            """Mock DNS resolution to simulate rebinding attack."""
            nonlocal call_count
            call_count += 1

            if call_count == 1:
                # First call (during validation): return public IP
                return [(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("8.8.8.8", 0))]
            # Second call (during httpx request): return localhost
            # This simulates the DNS rebinding attack
            # With DNS pinning, this should NOT be called
            return [(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("127.0.0.1", 0))]

        # Mock the network backend's connect_tcp to capture the actual IP being connected to
        async def mock_connect_tcp(self, host, port, **kwargs):
            """Capture the IP that's actually being connected to."""
            nonlocal connected_to_ip
            connected_to_ip = host
            # Return a mock stream with proper format (list of bytes)
            return httpcore.AsyncMockStream(
                [
                    b"HTTP/1.1 200 OK\r\n",
                    b"Content-Type: application/json\r\n",
                    b"Content-Length: 15\r\n",
                    b"\r\n",
                    b'{"status":"ok"}',
                ]
            )

        with (
            patch("socket.getaddrinfo", side_effect=mock_getaddrinfo),
            patch.dict(os.environ, {"LANGFLOW_SSRF_PROTECTION_ENABLED": "true"}),
            patch.object(httpcore.AnyIOBackend, "connect_tcp", mock_connect_tcp),
        ):
            # Execute the request
            result = await component.make_api_request()

            # Verify the request succeeded
            assert result is not None

            # CRITICAL CHECK: DNS should only be called once (during validation)
            # If called twice, DNS pinning failed and the attack succeeded
            assert call_count == 1, (
                f"DNS was called {call_count} times. Expected 1 (validation only). "
                "DNS pinning failed - the component is vulnerable to DNS rebinding!"
            )

            # Verify the connection was made to the pinned IP
            assert connected_to_ip is not None, "No TCP connection was made"
            assert connected_to_ip == "8.8.8.8", (
                f"Connection should be to pinned IP 8.8.8.8, but was to {connected_to_ip}"
            )

    @pytest.mark.asyncio
    async def test_dns_pinning_preserves_hostname_in_header(self, component):
        """Test that DNS pinning connects to pinned IP while preserving hostname for TLS.

        With network-level DNS pinning:
        - TCP connection goes to the pinned IP (93.184.216.34)
        - URL preserves the original hostname (rebinding.test)
        - This allows TLS SNI and certificate verification to work correctly

        This is important for:
        - Virtual hosting (multiple sites on same IP)
        - SNI (Server Name Indication) for HTTPS
        - Certificate verification (cert is for hostname, not IP)
        """
        connected_to_ip = None

        def mock_getaddrinfo(hostname, port, *args, **kwargs):
            """Mock DNS resolution."""
            return [(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("93.184.216.34", 0))]

        # Mock network backend to capture TCP connection
        async def mock_connect_tcp(self, host, port, **kwargs):
            """Capture the IP that TCP connects to."""
            nonlocal connected_to_ip
            connected_to_ip = host
            return httpcore.AsyncMockStream(
                [
                    b"HTTP/1.1 200 OK\r\n",
                    b"Content-Type: application/json\r\n",
                    b"Content-Length: 15\r\n",
                    b"\r\n",
                    b'{"status":"ok"}',
                ]
            )

        with (
            patch("socket.getaddrinfo", side_effect=mock_getaddrinfo),
            patch.dict(os.environ, {"LANGFLOW_SSRF_PROTECTION_ENABLED": "true"}),
            patch.object(httpcore.AnyIOBackend, "connect_tcp", mock_connect_tcp),
        ):
            result = await component.make_api_request()

            # Verify the request succeeded
            assert result is not None

            # Verify TCP connection was made to the pinned IP
            assert connected_to_ip is not None, "No TCP connection was made"
            assert connected_to_ip == "93.184.216.34", (
                f"TCP connection should be to pinned IP 93.184.216.34, but was to {connected_to_ip}"
            )

    @pytest.mark.asyncio
    async def test_dns_pinning_with_direct_ip_address(self, component):
        """Test that direct IP addresses work correctly (no DNS pinning needed)."""
        component.url_input = "http://93.184.216.34:8080/api"

        def mock_getaddrinfo(hostname, port, *args, **kwargs):
            """Mock DNS resolution - should not be called for direct IPs."""
            # For direct IPs, socket.getaddrinfo is still called but returns the same IP
            return [(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("93.184.216.34", 0))]

        mock_response = httpx.Response(200, json={"status": "ok"})

        with (
            patch("socket.getaddrinfo", side_effect=mock_getaddrinfo),
            patch.dict(os.environ, {"LANGFLOW_SSRF_PROTECTION_ENABLED": "true"}),
            patch("httpx.AsyncClient.request", return_value=mock_response) as mock_request,
        ):
            result = await component.make_api_request()

            # Verify the request succeeded
            assert isinstance(result, Data)
            assert mock_request.called

    @pytest.mark.asyncio
    async def test_dns_pinning_disabled_when_protection_disabled(self, component):
        """Test that DNS pinning is skipped when SSRF protection is disabled."""
        call_count = 0

        def mock_getaddrinfo(hostname, port, *args, **kwargs):
            """Mock DNS resolution."""
            nonlocal call_count
            call_count += 1
            return [(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("8.8.8.8", 0))]

        mock_response = httpx.Response(200, json={"status": "ok"})

        with (
            patch("socket.getaddrinfo", side_effect=mock_getaddrinfo),
            patch.dict(os.environ, {"LANGFLOW_SSRF_PROTECTION_ENABLED": "false"}),
            patch("httpx.AsyncClient.request", return_value=mock_response) as mock_request,
        ):
            result = await component.make_api_request()

            # Verify the request succeeded
            assert isinstance(result, Data)
            assert mock_request.called

            # When protection is disabled, DNS pinning is not used
            # So DNS might be called multiple times (once by httpx)
            # This is expected behavior when protection is off

    @pytest.mark.asyncio
    async def test_dns_pinning_blocks_private_ip_resolution(self, component):
        """Test that DNS pinning blocks hostnames that resolve to private IPs."""
        component.url_input = "http://internal.example.com/api"

        def mock_getaddrinfo(hostname, port, *args, **kwargs):
            """Mock DNS resolution to return private IP."""
            return [(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("192.168.1.1", 0))]

        with (
            patch("socket.getaddrinfo", side_effect=mock_getaddrinfo),
            patch.dict(os.environ, {"LANGFLOW_SSRF_PROTECTION_ENABLED": "true"}),
        ):
            # Should raise ValueError due to SSRF protection
            with pytest.raises(ValueError, match="SSRF Protection"):
                await component.make_api_request()

    @pytest.mark.asyncio
    async def test_dns_pinning_with_ipv6(self, component):
        """Test that DNS pinning works with IPv6 addresses.

        With network-level DNS pinning:
        - TCP connection goes to the IPv6 address
        - URL preserves the original hostname
        - IPv6 addresses don't need brackets in connect_tcp (only in URLs)
        """
        component.url_input = "http://ipv6.example.com/api"
        connected_to_ip = None

        def mock_getaddrinfo(hostname, port, *args, **kwargs):
            """Mock DNS resolution to return IPv6."""
            return [(socket.AF_INET6, socket.SOCK_STREAM, 6, "", ("2001:4860:4860::8888", 0))]

        # Mock network backend to capture TCP connection
        async def mock_connect_tcp(self, host, port, **kwargs):
            """Capture the IP that TCP connects to."""
            nonlocal connected_to_ip
            connected_to_ip = host
            return httpcore.AsyncMockStream(
                [
                    b"HTTP/1.1 200 OK\r\n",
                    b"Content-Type: application/json\r\n",
                    b"Content-Length: 15\r\n",
                    b"\r\n",
                    b'{"status":"ok"}',
                ]
            )

        with (
            patch("socket.getaddrinfo", side_effect=mock_getaddrinfo),
            patch.dict(os.environ, {"LANGFLOW_SSRF_PROTECTION_ENABLED": "true"}),
            patch.object(httpcore.AnyIOBackend, "connect_tcp", mock_connect_tcp),
        ):
            result = await component.make_api_request()

            # Verify the request succeeded
            assert isinstance(result, Data)

            # Verify TCP connection was made to the IPv6 address
            assert connected_to_ip is not None, "No TCP connection was made"
            assert connected_to_ip == "2001:4860:4860::8888", (
                f"TCP connection should be to IPv6 2001:4860:4860::8888, but was to {connected_to_ip}"
            )

    @pytest.mark.asyncio
    async def test_dns_pinning_with_allowlisted_host(self, component):
        """Test that allowlisted hosts bypass DNS pinning and preserve original hostname."""
        component.url_input = "http://internal.example.com:8080/api"  # Use a valid hostname format
        captured_request = None

        def mock_getaddrinfo(hostname, port, *args, **kwargs):
            """Mock DNS resolution."""
            return [(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("10.0.0.1", 0))]

        # Mock the transport's handle_async_request to capture the rewritten request
        async def mock_handle_async_request(_self, request):
            """Capture the request after transport rewrite."""
            nonlocal captured_request
            captured_request = request
            return httpx.Response(200, json={"status": "ok"}, request=request)

        with (
            patch("socket.getaddrinfo", side_effect=mock_getaddrinfo),
            patch.dict(
                os.environ,
                {
                    "LANGFLOW_SSRF_PROTECTION_ENABLED": "true",
                    "LANGFLOW_SSRF_ALLOWED_HOSTS": "internal.example.com,10.0.0.1",
                },
            ),
            # Patch get_allowed_hosts to return the allowlist directly
            patch(
                "lfx.utils.ssrf_protection.get_allowed_hosts",
                return_value=["internal.example.com", "10.0.0.1"],
            ),
            patch.object(httpx.AsyncHTTPTransport, "handle_async_request", mock_handle_async_request),
        ):
            result = await component.make_api_request()

            # Verify the request succeeded
            assert isinstance(result, Data)
            assert captured_request is not None, "Transport did not capture request"

            # Verify the original hostname is preserved (no DNS pinning for allowlisted hosts)
            url_str = str(captured_request.url)
            assert "internal.example.com" in url_str, f"Allowlisted host should preserve hostname: {url_str}"
            assert "10.0.0.1" not in url_str, f"Allowlisted host should not use IP: {url_str}"

    @pytest.mark.asyncio
    async def test_dns_pinning_with_multiple_ips_fallback(self, component):
        """Test that DNS pinning tries multiple IPs when first one fails (dual-stack/load balancing)."""
        component.url_input = "http://dual-stack.example.com/api"

        # Track which IPs were attempted
        attempted_ips = []

        def mock_getaddrinfo(host, port, family=0, type_=0, proto=0, flags=0):
            """Mock DNS resolution to return multiple IPs (IPv4 and IPv6)."""
            if host == "dual-stack.example.com":
                # Return both IPv4 and IPv6 addresses (dual-stack)
                return [
                    (socket.AF_INET, socket.SOCK_STREAM, 6, "", ("93.184.216.34", 0)),  # IPv4
                    (socket.AF_INET6, socket.SOCK_STREAM, 6, "", ("2606:2800:220:1:248:1893:25c8:1946", 0)),  # IPv6
                ]
            return [(socket.AF_INET, socket.SOCK_STREAM, 6, "", (host, 0))]

        async def mock_connect_tcp(self, host, port, **kwargs):
            """Mock connection that fails on first IP, succeeds on second."""
            nonlocal attempted_ips
            attempted_ips.append(host)

            # First IP fails (simulating IPv4 unreachable)
            if host == "93.184.216.34":
                msg = "Connection refused"
                raise OSError(msg)

            # Second IP succeeds (IPv6 works)
            if host == "2606:2800:220:1:248:1893:25c8:1946":
                return httpcore.AsyncMockStream(
                    [
                        b"HTTP/1.1 200 OK\r\n",
                        b"Content-Type: application/json\r\n",
                        b"Content-Length: 15\r\n",
                        b"\r\n",
                        b'{"status":"ok"}',
                    ]
                )

            msg = f"Unexpected host: {host}"
            raise OSError(msg)

        with (
            patch("socket.getaddrinfo", side_effect=mock_getaddrinfo),
            patch.dict(os.environ, {"LANGFLOW_SSRF_PROTECTION_ENABLED": "true"}),
            patch.object(httpcore.AnyIOBackend, "connect_tcp", mock_connect_tcp),
        ):
            # Execute the request
            result = await component.make_api_request()

            # Verify both IPs were attempted in order
            assert len(attempted_ips) == 2
            assert attempted_ips[0] == "93.184.216.34"  # IPv4 tried first
            assert attempted_ips[1] == "2606:2800:220:1:248:1893:25c8:1946"  # IPv6 tried second

            # Verify the result (should succeed with second IP)
            assert isinstance(result, Data)
            assert result.data is not None

    @pytest.mark.asyncio
    async def test_dns_pinning_applies_to_redirect_hops(self, component):
        """Test that DNS pinning is enforced on every redirect hop, not just the first.

        Simulates a redirect from a validated public host to a second host. With
        per-hop validation + pinning, both connections must go to the validated public
        IP - a rebinding attacker who flips the redirect target to 127.0.0.1 after
        validation can never cause a connection to the internal address.
        """
        component.url_input = "http://public.example.com/start"
        component.follow_redirects = True

        resolved = []

        def mock_getaddrinfo(host, *_args, **_kwargs):
            """Every host resolves to a public IP at validation time."""
            resolved.append(host)
            return [(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("8.8.8.8", 0))]

        connect_calls = []
        redirect_response = [
            b"HTTP/1.1 302 Found\r\n",
            b"Location: http://rebind.example.com/secret\r\n",
            b"Content-Length: 0\r\n",
            b"\r\n",
        ]
        final_response = [
            b"HTTP/1.1 200 OK\r\n",
            b"Content-Type: application/json\r\n",
            b"Content-Length: 15\r\n",
            b"\r\n",
            b'{"status":"ok"}',
        ]

        async def mock_connect_tcp(self, host, port, **kwargs):
            """Capture every IP connected to; first hop redirects, second succeeds."""
            connect_calls.append(host)
            stream = redirect_response if len(connect_calls) == 1 else final_response
            return httpcore.AsyncMockStream(list(stream))

        with (
            patch("socket.getaddrinfo", side_effect=mock_getaddrinfo),
            patch.dict(os.environ, {"LANGFLOW_SSRF_PROTECTION_ENABLED": "true"}),
            patch.object(httpcore.AnyIOBackend, "connect_tcp", mock_connect_tcp),
        ):
            result = await component.make_api_request()

        assert isinstance(result, Data)
        # Both the initial request and the redirect hop connected to the pinned public IP.
        assert connect_calls == ["8.8.8.8", "8.8.8.8"], connect_calls
        # The redirect target was resolved exactly once (validation only); httpx never
        # re-resolved it, so a rebinding flip after validation has no effect.
        assert resolved.count("rebind.example.com") == 1, resolved


class TestURLComponentDNSRebindingProtection:
    """Test DNS rebinding attack prevention in URL component.

    This test suite verifies that the URL component is protected against DNS rebinding
    attacks where an attacker controls a DNS server that returns different IPs on
    successive queries.
    """

    @pytest.mark.asyncio
    async def test_url_component_prevents_dns_rebinding_attack(self):
        """Test that URL component prevents DNS rebinding attacks via DNS pinning.

        This is a regression test for the SSRF vulnerability reported in the security disclosure.
        It verifies that only ONE DNS query occurs during the entire fetch operation.

        Attack scenario:
        1. First DNS query: attacker.test -> 1.1.1.1 (public IP, passes validation)
        2. Attacker changes DNS with TTL=0
        3. Second DNS query: attacker.test -> 127.0.0.1 (localhost, should be blocked)
        4. Without DNS pinning: fetch reaches localhost (SSRF bypass)
        5. With DNS pinning: fetch uses pinned 1.1.1.1 (attack prevented)
        """
        call_count = 0
        connected_to_ip = None

        def mock_getaddrinfo(_hostname, _port, *_args, **_kwargs):
            """Mock DNS resolution to simulate rebinding attack."""
            nonlocal call_count
            call_count += 1

            if call_count == 1:
                # First call (during validation): return public IP
                return [(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("1.1.1.1", 0))]
            # Second call (during httpx request): return localhost
            # This simulates the DNS rebinding attack
            # With DNS pinning, this should NOT be called
            return [(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("127.0.0.1", 0))]

        # Mock the network backend's connect_tcp to capture the actual IP being connected to
        async def mock_connect_tcp(self, host, port, **kwargs):
            """Capture the IP that's actually being connected to."""
            nonlocal connected_to_ip
            connected_to_ip = host
            # Return a mock stream with HTML content
            return httpcore.AsyncMockStream(
                [
                    b"HTTP/1.1 200 OK\r\n",
                    b"Content-Type: text/html\r\n",
                    b"Content-Length: 38\r\n",
                    b"\r\n",
                    b"<html><body>Test content</body></html>",
                ]
            )

        with (
            patch("socket.getaddrinfo", side_effect=mock_getaddrinfo),
            patch.dict(os.environ, {"LANGFLOW_SSRF_PROTECTION_ENABLED": "true"}),
            patch.object(httpcore.AnyIOBackend, "connect_tcp", mock_connect_tcp),
        ):
            # Create URL component
            component = URLComponent()
            component.urls = ["http://attacker.test/"]
            component.max_depth = 1
            component.format = "Text"
            component.headers = [{"key": "User-Agent", "value": "Test"}]
            component.timeout = 30
            component.prevent_outside = True
            component.use_async = True
            component.filter_text_html = False
            component.continue_on_failure = False
            component.check_response_status = False
            component.autoset_encoding = True

            # Execute the component
            result = await component.fetch_url_contents()

            # Verify results
            assert len(result) == 1
            assert "Test content" in result[0]["text"]

            # CRITICAL VERIFICATION: Only ONE DNS query should have occurred
            assert call_count == 1, (
                f"DNS rebinding attack NOT prevented! "
                f"Expected 1 DNS query (during validation), but {call_count} occurred. "
                f"The URL component is vulnerable to DNS rebinding attacks."
            )

            # Verify the connection was made to the pinned IP
            assert connected_to_ip is not None, "No TCP connection was made"
            assert connected_to_ip == "1.1.1.1", (
                f"Connection should be to pinned IP 1.1.1.1, but was to {connected_to_ip}. DNS pinning failed!"
            )

    @pytest.mark.asyncio
    async def test_url_component_blocks_localhost_on_first_query(self):
        """Test that localhost is blocked even on the first DNS query."""

        def mock_getaddrinfo_localhost(hostname, port, *args, **kwargs):
            """Return localhost immediately."""
            return [
                (
                    socket.AF_INET,
                    socket.SOCK_STREAM,
                    6,
                    "",
                    ("127.0.0.1", port or 0),
                )
            ]

        with (
            patch("socket.getaddrinfo", side_effect=mock_getaddrinfo_localhost),
            patch.dict(os.environ, {"LANGFLOW_SSRF_PROTECTION_ENABLED": "true"}),
        ):
            component = URLComponent()
            component.urls = ["http://localhost-test.example/"]
            component.max_depth = 1
            component.format = "Text"
            component.headers = []
            component.timeout = 30

            # Should raise ValueError due to SSRF protection
            with pytest.raises(ValueError, match="SSRF Protection"):
                await component.fetch_url_contents()

    @pytest.mark.asyncio
    async def test_url_component_blocks_aws_metadata_on_first_query(self):
        """Test that AWS metadata endpoint is blocked even on the first DNS query."""

        def mock_getaddrinfo_metadata(hostname, port, *args, **kwargs):
            """Return AWS metadata IP immediately."""
            return [
                (
                    socket.AF_INET,
                    socket.SOCK_STREAM,
                    6,
                    "",
                    ("169.254.169.254", port or 0),
                )
            ]

        with (
            patch("socket.getaddrinfo", side_effect=mock_getaddrinfo_metadata),
            patch.dict(os.environ, {"LANGFLOW_SSRF_PROTECTION_ENABLED": "true"}),
        ):
            component = URLComponent()
            component.urls = ["http://metadata.example/"]
            component.max_depth = 1
            component.format = "Text"
            component.headers = []
            component.timeout = 30

            # Should raise ValueError due to SSRF protection
            with pytest.raises(ValueError, match="SSRF Protection"):
                await component.fetch_url_contents()

    @pytest.mark.asyncio
    async def test_url_component_revalidates_discovered_links(self):
        """Test that each discovered link during recursive crawling is re-validated.

        This ensures that even if the initial URL is safe, any links discovered
        during crawling are also validated before being fetched.
        """
        call_count = 0
        validated_hosts = []

        def mock_getaddrinfo(hostname, port, *args, **kwargs):
            """Track which hosts are validated."""
            nonlocal call_count
            call_count += 1
            validated_hosts.append(hostname)

            # First host (safe.com) resolves to public IP
            if hostname == "safe.com":
                return [(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("1.1.1.1", 0))]
            # Second host (evil.com) resolves to localhost - should be blocked
            if hostname == "evil.com":
                return [(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("127.0.0.1", 0))]
            return [(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("8.8.8.8", 0))]

        async def mock_connect_tcp(self, host, port, **kwargs):
            """Return HTML with a link to evil.com."""
            return httpcore.AsyncMockStream(
                [
                    b"HTTP/1.1 200 OK\r\n",
                    b"Content-Type: text/html\r\n",
                    b"Content-Length: 61\r\n",
                    b"\r\n",
                    b'<html><body><a href="http://evil.com/">Link</a></body></html>',
                ]
            )

        with (
            patch("socket.getaddrinfo", side_effect=mock_getaddrinfo),
            patch.dict(os.environ, {"LANGFLOW_SSRF_PROTECTION_ENABLED": "true"}),
            patch.object(httpcore.AnyIOBackend, "connect_tcp", mock_connect_tcp),
        ):
            component = URLComponent()
            component.urls = ["http://safe.com/"]
            component.max_depth = 2  # Allow following links
            component.format = "Text"
            component.headers = []
            component.timeout = 30
            component.prevent_outside = False  # Allow external links
            component.continue_on_failure = True  # Continue even if evil.com is blocked

            # Execute - should fetch safe.com but block evil.com
            result = await component.fetch_url_contents()

            # Verify safe.com was fetched
            assert len(result) >= 1

            # Verify both hosts were validated
            assert "safe.com" in validated_hosts, "Initial URL should be validated"
            assert "evil.com" in validated_hosts, "Discovered link should be validated"

            # Verify evil.com was blocked (only safe.com content returned)
            assert len(result) == 1, "evil.com should have been blocked, only safe.com content returned"

    @pytest.mark.asyncio
    async def test_url_component_blocks_mixed_safe_unsafe_dns_answers(self):
        """Test that hostnames resolving to both safe and unsafe IPs are blocked.

        This prevents attacks where a hostname resolves to multiple IPs including
        both public (safe) and private/localhost (unsafe) addresses. The entire
        hostname must be blocked if ANY resolved IP is unsafe.
        """

        def mock_getaddrinfo(hostname, port, *args, **kwargs):
            """Return both safe and unsafe IPs for the same hostname."""
            return [
                # First IP: safe public address
                (socket.AF_INET, socket.SOCK_STREAM, 6, "", ("8.8.8.8", 0)),
                # Second IP: unsafe localhost
                (socket.AF_INET, socket.SOCK_STREAM, 6, "", ("127.0.0.1", 0)),
            ]

        with (
            patch("socket.getaddrinfo", side_effect=mock_getaddrinfo),
            patch.dict(os.environ, {"LANGFLOW_SSRF_PROTECTION_ENABLED": "true"}),
        ):
            component = URLComponent()
            component.urls = ["http://mixed-dns.example/"]
            component.max_depth = 1
            component.format = "Text"
            component.headers = []
            component.timeout = 30

            # Should raise ValueError due to SSRF protection blocking the unsafe IP
            with pytest.raises(ValueError, match=r"SSRF Protection.*127\.0\.0\.1"):
                await component.fetch_url_contents()

    @pytest.mark.asyncio
    async def test_url_component_blocks_ipv4_mapped_ipv6_localhost(self):
        """Test that IPv4-mapped IPv6 localhost addresses are blocked.

        IPv4-mapped IPv6 addresses like ::ffff:127.0.0.1 are a common SSRF bypass
        technique. The protection must recognize and block these.
        """

        def mock_getaddrinfo(hostname, port, *args, **kwargs):
            """Return IPv4-mapped IPv6 localhost."""
            return [
                (
                    socket.AF_INET6,
                    socket.SOCK_STREAM,
                    6,
                    "",
                    ("::ffff:127.0.0.1", 0),
                )
            ]

        with (
            patch("socket.getaddrinfo", side_effect=mock_getaddrinfo),
            patch.dict(os.environ, {"LANGFLOW_SSRF_PROTECTION_ENABLED": "true"}),
        ):
            component = URLComponent()
            component.urls = ["http://ipv4-mapped.example/"]
            component.max_depth = 1
            component.format = "Text"
            component.headers = []
            component.timeout = 30

            # Should raise ValueError due to SSRF protection
            with pytest.raises(ValueError, match="SSRF Protection"):
                await component.fetch_url_contents()

    @pytest.mark.asyncio
    async def test_url_component_blocks_ipv4_mapped_ipv6_private_network(self):
        """Test that IPv4-mapped IPv6 private network addresses are blocked.

        Tests ::ffff:192.168.1.1 (IPv4-mapped private network).
        """

        def mock_getaddrinfo(hostname, port, *args, **kwargs):
            """Return IPv4-mapped IPv6 private network address."""
            return [
                (
                    socket.AF_INET6,
                    socket.SOCK_STREAM,
                    6,
                    "",
                    ("::ffff:192.168.1.1", 0),
                )
            ]

        with (
            patch("socket.getaddrinfo", side_effect=mock_getaddrinfo),
            patch.dict(os.environ, {"LANGFLOW_SSRF_PROTECTION_ENABLED": "true"}),
        ):
            component = URLComponent()
            component.urls = ["http://ipv4-mapped-private.example/"]
            component.max_depth = 1
            component.format = "Text"
            component.headers = []
            component.timeout = 30

            # Should raise ValueError due to SSRF protection
            with pytest.raises(ValueError, match="SSRF Protection"):
                await component.fetch_url_contents()

    @pytest.mark.asyncio
    async def test_url_component_follows_redirects_with_per_hop_pinning(self):
        """Test that the URL component follows redirects and re-validates/pins every hop.

        Regression test: the component previously sent every request with
        follow_redirects=False, so a site that 301s to its canonical URL (http->https or
        www normalization - extremely common) returned the redirect stub instead of the
        page. Redirects are now followed, and each hop is independently DNS-pinned so a
        rebinding attacker who flips the redirect target after validation can never cause
        a connection to an unvalidated address.
        """
        resolved = []

        def mock_getaddrinfo(host, *_args, **_kwargs):
            """Every host resolves to the same public IP at validation time."""
            resolved.append(host)
            return [(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("8.8.8.8", 0))]

        connect_hosts = []
        redirect_response = [
            b"HTTP/1.1 301 Moved Permanently\r\n",
            b"Location: https://final.test/\r\n",
            b"Content-Length: 0\r\n",
            b"\r\n",
        ]
        final_response = [
            b"HTTP/1.1 200 OK\r\n",
            b"Content-Type: text/html\r\n",
            b"Content-Length: 39\r\n",
            b"\r\n",
            b"<html><body>Final content</body></html>",
        ]

        async def mock_connect_tcp(self, host, port, **kwargs):
            """Capture every IP connected to; first hop redirects, second succeeds."""
            connect_hosts.append(host)
            stream = redirect_response if len(connect_hosts) == 1 else final_response
            return httpcore.AsyncMockStream(list(stream))

        with (
            patch("socket.getaddrinfo", side_effect=mock_getaddrinfo),
            patch.dict(os.environ, {"LANGFLOW_SSRF_PROTECTION_ENABLED": "true"}),
            patch.object(httpcore.AnyIOBackend, "connect_tcp", mock_connect_tcp),
        ):
            component = URLComponent()
            component.urls = ["http://redirect.test/"]
            component.max_depth = 1
            component.format = "Text"
            component.headers = [{"key": "User-Agent", "value": "Test"}]
            component.timeout = 30
            component.prevent_outside = True
            component.use_async = True
            component.follow_redirects = True
            component.filter_text_html = False
            component.continue_on_failure = False
            component.check_response_status = False
            component.autoset_encoding = True

            result = await component.fetch_url_contents()

        # The redirect was followed through to the final page's content.
        assert len(result) == 1
        assert "Final content" in result[0]["text"]
        # Both the initial request and the redirect hop connected to the pinned public IP.
        assert connect_hosts == ["8.8.8.8", "8.8.8.8"], connect_hosts
        # The redirect target was resolved exactly once (validation only); httpx never
        # re-resolved it, so a rebinding flip after validation has no effect.
        assert resolved.count("final.test") == 1, resolved

    @pytest.mark.asyncio
    async def test_url_component_blocks_redirect_to_internal_ip(self):
        """Test that a redirect pointing at an internal address is blocked before connecting.

        A validated public host redirects to a host that resolves to localhost. The hop
        is re-validated with the SSRF denylist, so the component raises rather than
        fetching the internal resource, and never opens a connection to it.
        """
        resolved = []

        def mock_getaddrinfo(host, *_args, **_kwargs):
            """The public redirector resolves to a public IP; the target to localhost."""
            resolved.append(host)
            if host == "internal.test":
                return [(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("127.0.0.1", 0))]
            return [(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("8.8.8.8", 0))]

        connect_hosts = []
        redirect_response = [
            b"HTTP/1.1 302 Found\r\n",
            b"Location: http://internal.test/secret\r\n",
            b"Content-Length: 0\r\n",
            b"\r\n",
        ]

        async def mock_connect_tcp(self, host, port, **kwargs):
            """First (and only) connection returns a redirect to the internal host."""
            connect_hosts.append(host)
            return httpcore.AsyncMockStream(list(redirect_response))

        with (
            patch("socket.getaddrinfo", side_effect=mock_getaddrinfo),
            patch.dict(os.environ, {"LANGFLOW_SSRF_PROTECTION_ENABLED": "true"}),
            patch.object(httpcore.AnyIOBackend, "connect_tcp", mock_connect_tcp),
        ):
            component = URLComponent()
            component.urls = ["http://safe-redirector.test/"]
            component.max_depth = 1
            component.format = "Text"
            component.headers = [{"key": "User-Agent", "value": "Test"}]
            component.timeout = 30
            component.prevent_outside = True
            component.use_async = True
            component.follow_redirects = True
            component.filter_text_html = False
            component.continue_on_failure = False
            component.check_response_status = False
            component.autoset_encoding = True

            # The redirect to a localhost-resolving host must be blocked by SSRF protection.
            with pytest.raises(ValueError, match="SSRF Protection"):
                await component.fetch_url_contents()

        # Only the first (public) hop ever connected; the internal target was never reached.
        assert connect_hosts == ["8.8.8.8"], connect_hosts
        assert "internal.test" in resolved, resolved

    @pytest.mark.asyncio
    async def test_url_component_crawls_links_from_redirected_base(self):
        """Test that depth>1 crawls resolve links against the final post-redirect URL.

        Regression test: after following http://redirect.test/ -> https://final.test/,
        relative links were still resolved against the original URL (sending /about to
        http://redirect.test/about) and prevent_outside compared link domains against
        the pre-redirect host, skipping same-site absolute links. Links must be resolved
        against the URL the content actually came from, and the canonical URL itself is
        marked visited so links back to it are not re-crawled.
        """
        resolved = []

        def mock_getaddrinfo(host, *_args, **_kwargs):
            """Every host resolves to the same public IP at validation time."""
            resolved.append(host)
            return [(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("8.8.8.8", 0))]

        def http_ok(body: bytes) -> list[bytes]:
            return [
                b"HTTP/1.1 200 OK\r\n",
                b"Content-Type: text/html\r\n",
                b"Content-Length: " + str(len(body)).encode() + b"\r\n",
                b"\r\n",
                body,
            ]

        responses = [
            [
                b"HTTP/1.1 301 Moved Permanently\r\n",
                b"Location: https://final.test/\r\n",
                b"Content-Length: 0\r\n",
                b"\r\n",
            ],
            http_ok(
                b"<html><body>Final page"
                b'<a href="/about">About</a>'
                b'<a href="https://final.test/contact">Contact</a>'
                b'<a href="https://final.test/">Home</a>'
                b"</body></html>"
            ),
            http_ok(b"<html><body>About page</body></html>"),
            http_ok(b"<html><body>Contact page</body></html>"),
        ]

        connect_hosts = []

        async def mock_connect_tcp(self, host, port, **kwargs):
            """Serve the redirect, the final page, then one page per crawled link."""
            connect_hosts.append(host)
            return httpcore.AsyncMockStream(list(responses[len(connect_hosts) - 1]))

        with (
            patch("socket.getaddrinfo", side_effect=mock_getaddrinfo),
            patch.dict(os.environ, {"LANGFLOW_SSRF_PROTECTION_ENABLED": "true"}),
            patch.object(httpcore.AnyIOBackend, "connect_tcp", mock_connect_tcp),
        ):
            component = URLComponent()
            component.urls = ["http://redirect.test/"]
            component.max_depth = 2
            component.format = "Text"
            component.headers = [{"key": "User-Agent", "value": "Test"}]
            component.timeout = 30
            component.prevent_outside = True
            component.use_async = True
            component.follow_redirects = True
            component.filter_text_html = False
            component.continue_on_failure = False
            component.check_response_status = False
            component.autoset_encoding = True

            result = await component.fetch_url_contents()

        # The final page and both same-site links (relative and absolute) were crawled.
        assert [doc["url"] for doc in result] == [
            "https://final.test/",
            "https://final.test/about",
            "https://final.test/contact",
        ]
        assert "About page" in result[1]["text"]
        # The back-link to the canonical URL was already visited - exactly 4 connections
        # (initial + redirect hop + 2 links), all to the pinned public IP.
        assert connect_hosts == ["8.8.8.8"] * 4, connect_hosts
        # Links validated against the post-redirect host; the pre-redirect host was only
        # resolved once, for the initial URL validation.
        assert resolved.count("redirect.test") == 1, resolved
        assert resolved.count("final.test") == 3, resolved


class TestAPIRequestDNSRebindingEdgeCases:
    """Additional edge case tests for API Request component DNS rebinding protection."""

    @pytest.fixture
    def component(self):
        """Create a basic API request component."""
        return APIRequestComponent(
            url_input="http://test.example:8080/api",
            method="GET",
            headers=[],
            body=[],
            timeout=30,
            follow_redirects=False,
            save_to_file=False,
            include_httpx_metadata=False,
            mode="URL",
            curl_input="",
            query_params={},
        )

    @pytest.mark.asyncio
    async def test_api_request_blocks_mixed_safe_unsafe_dns_answers(self, component):
        """Test that hostnames resolving to both safe and unsafe IPs are blocked."""

        def mock_getaddrinfo(hostname, port, *args, **kwargs):
            """Return both safe and unsafe IPs."""
            return [
                (socket.AF_INET, socket.SOCK_STREAM, 6, "", ("8.8.8.8", 0)),
                (socket.AF_INET, socket.SOCK_STREAM, 6, "", ("127.0.0.1", 0)),
            ]

        with (
            patch("socket.getaddrinfo", side_effect=mock_getaddrinfo),
            patch.dict(os.environ, {"LANGFLOW_SSRF_PROTECTION_ENABLED": "true"}),
        ):
            with pytest.raises(ValueError, match=r"SSRF Protection.*127\.0\.0\.1"):
                await component.make_api_request()

    @pytest.mark.asyncio
    async def test_api_request_blocks_ipv4_mapped_ipv6_localhost(self, component):
        """Test that IPv4-mapped IPv6 localhost is blocked."""

        def mock_getaddrinfo(hostname, port, *args, **kwargs):
            """Return IPv4-mapped IPv6 localhost."""
            return [(socket.AF_INET6, socket.SOCK_STREAM, 6, "", ("::ffff:127.0.0.1", 0))]

        with (
            patch("socket.getaddrinfo", side_effect=mock_getaddrinfo),
            patch.dict(os.environ, {"LANGFLOW_SSRF_PROTECTION_ENABLED": "true"}),
        ):
            with pytest.raises(ValueError, match="SSRF Protection"):
                await component.make_api_request()

    @pytest.mark.asyncio
    async def test_api_request_blocks_ipv4_mapped_ipv6_metadata(self, component):
        """Test that IPv4-mapped IPv6 AWS metadata is blocked."""

        def mock_getaddrinfo(hostname, port, *args, **kwargs):
            """Return IPv4-mapped IPv6 AWS metadata address."""
            return [(socket.AF_INET6, socket.SOCK_STREAM, 6, "", ("::ffff:169.254.169.254", 0))]

        with (
            patch("socket.getaddrinfo", side_effect=mock_getaddrinfo),
            patch.dict(os.environ, {"LANGFLOW_SSRF_PROTECTION_ENABLED": "true"}),
        ):
            with pytest.raises(ValueError, match="SSRF Protection"):
                await component.make_api_request()