Skyvern/tests/unit/test_cdp_download_interceptor.py

"""Unit tests for CDPDownloadInterceptor pure functions and proxy auth handling."""

from pathlib import Path
from unittest.mock import AsyncMock, MagicMock

import pytest

from skyvern.webeye.cdp_download_interceptor import CDPDownloadInterceptor, extract_filename, is_download_response


class TestIsDownloadResponse:
    """Tests for is_download_response()."""

    def test_attachment_header(self) -> None:
        headers = {"content-disposition": 'attachment; filename="report.csv"', "content-type": "text/csv"}
        assert is_download_response(headers, 200) is True

    def test_attachment_header_case_insensitive(self) -> None:
        headers = {"content-disposition": 'Attachment; filename="report.csv"', "content-type": "text/csv"}
        assert is_download_response(headers, 200) is True

    def test_download_mime_type_pdf(self) -> None:
        headers = {"content-type": "application/pdf"}
        assert is_download_response(headers, 200) is True

    def test_download_mime_type_zip(self) -> None:
        headers = {"content-type": "application/zip"}
        assert is_download_response(headers, 200) is True

    def test_download_mime_type_xlsx(self) -> None:
        headers = {
            "content-type": "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
        }
        assert is_download_response(headers, 200) is True

    def test_download_mime_type_octet_stream(self) -> None:
        headers = {"content-type": "application/octet-stream"}
        assert is_download_response(headers, 200) is True

    def test_download_mime_type_with_charset(self) -> None:
        headers = {"content-type": "application/pdf; charset=utf-8"}
        assert is_download_response(headers, 200) is True

    def test_html_not_download(self) -> None:
        headers = {"content-type": "text/html"}
        assert is_download_response(headers, 200) is False

    def test_json_not_download(self) -> None:
        headers = {"content-type": "application/json"}
        assert is_download_response(headers, 200) is False

    def test_json_with_attachment_not_download(self) -> None:
        """JSON responses with Content-Disposition: attachment should NOT be treated as downloads."""
        headers = {
            "content-disposition": "attachment",
            "content-type": "application/json",
        }
        assert is_download_response(headers, 200) is False

    def test_xml_not_download(self) -> None:
        headers = {"content-type": "application/xml"}
        assert is_download_response(headers, 200) is False

    def test_grpc_not_download(self) -> None:
        headers = {"content-type": "application/grpc"}
        assert is_download_response(headers, 200) is False

    def test_empty_headers_not_download(self) -> None:
        assert is_download_response({}, 200) is False

    # Resource type filtering — XHR/Fetch require BOTH attachment AND download MIME
    def test_xhr_attachment_and_download_mime_is_download(self) -> None:
        """XHR with attachment + download MIME type should be treated as download."""
        headers = {"content-disposition": "attachment", "content-type": "application/octet-stream"}
        assert is_download_response(headers, 200, resource_type="XHR") is True

    def test_fetch_attachment_and_download_mime_is_download(self) -> None:
        """Fetch with attachment + download MIME type should be treated as download."""
        headers = {"content-disposition": "attachment", "content-type": "application/pdf"}
        assert is_download_response(headers, 200, resource_type="Fetch") is True

    def test_xhr_attachment_pdf_is_download(self) -> None:
        """Real-world case: XHR download with attachment header and PDF content-type."""
        headers = {
            "content-disposition": "attachment; filename=Invoice_12345.pdf; filename*=UTF-8''Invoice_12345.pdf",
            "content-type": "application/pdf",
        }
        assert is_download_response(headers, 200, resource_type="XHR") is True

    # XHR/Fetch with attachment but non-download MIME should NOT be download
    def test_xhr_attachment_text_plain_not_download(self) -> None:
        """Google-style XHR: text/plain + attachment should NOT be treated as download."""
        headers = {"content-disposition": 'attachment; filename="f.txt"', "content-type": "text/plain; charset=UTF-8"}
        assert is_download_response(headers, 200, resource_type="XHR") is False

    def test_xhr_attachment_text_html_not_download(self) -> None:
        """XHR with text/html + attachment should NOT be treated as download."""
        headers = {"content-disposition": "attachment", "content-type": "text/html"}
        assert is_download_response(headers, 200, resource_type="XHR") is False

    def test_fetch_attachment_only_not_download(self) -> None:
        """Fetch with attachment but no download MIME type should NOT be download."""
        headers = {"content-disposition": "attachment"}
        assert is_download_response(headers, 200, resource_type="Fetch") is False

    def test_xhr_attachment_csv_not_download(self) -> None:
        """Known limitation: CSV via XHR is not detected because text/csv is not in DOWNLOAD_MIME_TYPES."""
        headers = {"content-disposition": "attachment", "content-type": "text/csv"}
        assert is_download_response(headers, 200, resource_type="XHR") is False

    # XHR/Fetch without attachment header should NOT be download (MIME-only false positive)
    def test_xhr_mime_only_not_download(self) -> None:
        """XHR with download MIME type but no attachment header should NOT be treated as download."""
        headers = {"content-type": "application/pdf"}
        assert is_download_response(headers, 200, resource_type="XHR") is False

    def test_fetch_mime_only_not_download(self) -> None:
        """Fetch with download MIME type but no attachment header should NOT be treated as download."""
        headers = {"content-type": "application/octet-stream"}
        assert is_download_response(headers, 200, resource_type="Fetch") is False

    # XHR/Fetch with attachment but API content-type should still be filtered
    def test_xhr_json_attachment_not_download(self) -> None:
        """XHR with JSON content-type and attachment header should NOT be download."""
        headers = {"content-disposition": "attachment", "content-type": "application/json"}
        assert is_download_response(headers, 200, resource_type="XHR") is False

    def test_font_resource_type_not_download(self) -> None:
        headers = {"content-type": "application/octet-stream"}
        assert is_download_response(headers, 200, resource_type="Font") is False

    def test_stylesheet_resource_type_not_download(self) -> None:
        headers = {"content-type": "application/octet-stream"}
        assert is_download_response(headers, 200, resource_type="Stylesheet") is False

    def test_script_resource_type_not_download(self) -> None:
        headers = {"content-type": "application/octet-stream"}
        assert is_download_response(headers, 200, resource_type="Script") is False

    def test_image_resource_type_not_download(self) -> None:
        headers = {"content-type": "application/octet-stream"}
        assert is_download_response(headers, 200, resource_type="Image") is False

    def test_document_resource_type_is_download(self) -> None:
        """Document resource type (link click) should allow download detection."""
        headers = {"content-disposition": "attachment", "content-type": "application/pdf"}
        assert is_download_response(headers, 200, resource_type="Document") is True

    def test_empty_resource_type_is_download(self) -> None:
        headers = {"content-type": "application/pdf"}
        assert is_download_response(headers, 200, resource_type="") is True

    def test_error_status_code_not_download(self) -> None:
        headers = {"content-disposition": "attachment", "content-type": "application/pdf"}
        assert is_download_response(headers, 404) is False

    def test_server_error_not_download(self) -> None:
        headers = {"content-type": "application/octet-stream"}
        assert is_download_response(headers, 500) is False


class TestExtractFilename:
    """Tests for extract_filename().

    extract_filename returns an empty string when no filename can be determined —
    the caller (_resolve_save_path) is responsible for generating a fallback name.
    """

    def test_rfc5987_filename_star(self) -> None:
        headers = {"content-disposition": "attachment; filename*=UTF-8''my%20report%282024%29.pdf"}
        result = extract_filename(headers, "https://example.com/download")
        assert result == "my report(2024).pdf"

    def test_regular_filename(self) -> None:
        headers = {"content-disposition": 'attachment; filename="report.csv"'}
        result = extract_filename(headers, "https://example.com/download")
        assert result == "report.csv"

    def test_unquoted_filename(self) -> None:
        headers = {"content-disposition": "attachment; filename=report.csv"}
        result = extract_filename(headers, "https://example.com/download")
        assert result == "report.csv"

    def test_filename_star_takes_priority(self) -> None:
        headers = {
            "content-disposition": "attachment; filename=\"fallback.csv\"; filename*=UTF-8''preferred.csv",
        }
        result = extract_filename(headers, "https://example.com/download")
        assert result == "preferred.csv"

    def test_url_path_fallback(self) -> None:
        headers: dict[str, str] = {}
        result = extract_filename(headers, "https://example.com/files/document.pdf")
        assert result == "document.pdf"

    def test_url_path_with_encoded_chars(self) -> None:
        headers: dict[str, str] = {}
        result = extract_filename(headers, "https://example.com/files/my%20report.xlsx")
        assert result == "my report.xlsx"

    def test_url_path_no_extension_returns_empty(self) -> None:
        """No extension in URL path and no Content-Disposition — returns empty string."""
        headers: dict[str, str] = {}
        result = extract_filename(headers, "https://example.com/download")
        assert result == ""

    def test_no_headers_no_url_returns_empty(self) -> None:
        """Completely empty inputs — returns empty string for _resolve_save_path to handle."""
        result = extract_filename({}, "https://example.com/api/export")
        assert result == ""

    def test_empty_content_disposition(self) -> None:
        headers = {"content-disposition": ""}
        result = extract_filename(headers, "https://example.com/files/data.csv")
        assert result == "data.csv"

    def test_content_disposition_inline(self) -> None:
        """inline disposition without filename should fall back to URL."""
        headers = {"content-disposition": "inline"}
        result = extract_filename(headers, "https://example.com/files/report.pdf")
        assert result == "report.pdf"

    def test_path_traversal_returned_raw(self) -> None:
        """extract_filename returns raw name; sanitization is done in _resolve_save_path."""
        headers = {"content-disposition": 'attachment; filename="../../etc/cron.d/evil"'}
        result = extract_filename(headers, "https://example.com/download")
        assert result == "../../etc/cron.d/evil"


class TestResolveSavePath:
    """Tests for CDPDownloadInterceptor._resolve_save_path()."""

    def _make_interceptor(self, tmp_path: Path) -> CDPDownloadInterceptor:
        interceptor = CDPDownloadInterceptor(output_dir=str(tmp_path))
        return interceptor

    def test_normal_filename(self, tmp_path: Path) -> None:
        interceptor = self._make_interceptor(tmp_path)
        save_path, filename = interceptor._resolve_save_path("report.pdf")
        assert filename == "report.pdf"
        assert save_path == tmp_path / "report.pdf"

    def test_empty_filename_gets_uuid_fallback(self, tmp_path: Path) -> None:
        """Empty filename should generate a download_{uuid} fallback."""
        interceptor = self._make_interceptor(tmp_path)
        save_path, filename = interceptor._resolve_save_path("")
        assert filename.startswith("download_")
        assert len(filename) > len("download_")
        assert save_path == tmp_path / filename

    def test_default_param_empty_string(self, tmp_path: Path) -> None:
        """Calling without arguments should also trigger fallback."""
        interceptor = self._make_interceptor(tmp_path)
        _, filename = interceptor._resolve_save_path()
        assert filename.startswith("download_")

    def test_path_traversal_sanitized(self, tmp_path: Path) -> None:
        """Path traversal components should be stripped — only the final name is kept."""
        interceptor = self._make_interceptor(tmp_path)
        save_path, filename = interceptor._resolve_save_path("../../etc/cron.d/evil")
        assert filename == "evil"
        assert save_path == tmp_path / "evil"

    def test_increments_download_index(self, tmp_path: Path) -> None:
        interceptor = self._make_interceptor(tmp_path)
        assert interceptor._download_index == 0
        interceptor._resolve_save_path("a.pdf")
        assert interceptor._download_index == 1
        interceptor._resolve_save_path("b.pdf")
        assert interceptor._download_index == 2

    def test_collision_warning_logged(self, tmp_path: Path) -> None:
        """Existing file with the same name should warn but still return the path."""
        interceptor = self._make_interceptor(tmp_path)
        # Create a file that will collide
        (tmp_path / "report.pdf").write_bytes(b"existing")
        save_path, filename = interceptor._resolve_save_path("report.pdf")
        assert filename == "report.pdf"
        assert save_path == tmp_path / "report.pdf"

    def test_creates_output_dir_if_missing(self, tmp_path: Path) -> None:
        nested = tmp_path / "sub" / "dir"
        interceptor = CDPDownloadInterceptor(output_dir=str(nested))
        save_path, _ = interceptor._resolve_save_path("file.txt")
        assert nested.exists()
        assert save_path.parent == nested


class TestCDPDownloadInterceptorProxyAuth:
    """Tests for CDP proxy authentication handling (Fetch.authRequired + continueWithAuth)."""

    def _make_interceptor(
        self,
        proxy_username: str | None = None,
        proxy_password: str | None = None,
    ) -> CDPDownloadInterceptor:
        return CDPDownloadInterceptor(
            output_dir="/tmp/test_downloads",
            proxy_username=proxy_username,
            proxy_password=proxy_password,
        )

    def _make_cdp_session(self) -> MagicMock:
        session = MagicMock()
        session.send = AsyncMock()
        return session

    @pytest.mark.asyncio
    async def test_proxy_auth_provides_credentials(self) -> None:
        """Proxy 407 challenge should respond with ProvideCredentials when credentials are available."""
        interceptor = self._make_interceptor(proxy_username="user1", proxy_password="pass1")
        cdp_session = self._make_cdp_session()

        event = {
            "requestId": "req-1",
            "authChallenge": {"source": "Proxy", "origin": "http://proxy.example.com"},
            "request": {"url": "https://example.com/page"},
        }

        await interceptor._handle_auth_required(event, cdp_session)

        cdp_session.send.assert_called_once_with(
            "Fetch.continueWithAuth",
            {
                "requestId": "req-1",
                "authChallengeResponse": {
                    "response": "ProvideCredentials",
                    "username": "user1",
                    "password": "pass1",
                },
            },
        )

    @pytest.mark.asyncio
    async def test_non_proxy_auth_cancels(self) -> None:
        """Non-proxy auth challenges (e.g., HTTP Basic from origin) should be cancelled."""
        interceptor = self._make_interceptor(proxy_username="user1", proxy_password="pass1")
        cdp_session = self._make_cdp_session()

        event = {
            "requestId": "req-2",
            "authChallenge": {"source": "Server", "origin": "https://example.com"},
            "request": {"url": "https://example.com/protected"},
        }

        await interceptor._handle_auth_required(event, cdp_session)

        cdp_session.send.assert_called_once_with(
            "Fetch.continueWithAuth",
            {
                "requestId": "req-2",
                "authChallengeResponse": {"response": "CancelAuth"},
            },
        )

    @pytest.mark.asyncio
    async def test_no_credentials_cancels_proxy_auth(self) -> None:
        """Proxy auth challenge without credentials should be cancelled."""
        interceptor = self._make_interceptor()  # No credentials
        cdp_session = self._make_cdp_session()

        event = {
            "requestId": "req-3",
            "authChallenge": {"source": "Proxy", "origin": "http://proxy.example.com"},
            "request": {"url": "https://example.com/page"},
        }

        await interceptor._handle_auth_required(event, cdp_session)

        cdp_session.send.assert_called_once_with(
            "Fetch.continueWithAuth",
            {
                "requestId": "req-3",
                "authChallengeResponse": {"response": "CancelAuth"},
            },
        )

    @pytest.mark.asyncio
    async def test_partial_credentials_cancels(self) -> None:
        """Proxy auth with only username (no password) should cancel."""
        interceptor = self._make_interceptor(proxy_username="user1")
        cdp_session = self._make_cdp_session()

        event = {
            "requestId": "req-4",
            "authChallenge": {"source": "Proxy", "origin": "http://proxy.example.com"},
            "request": {"url": "https://example.com/page"},
        }

        await interceptor._handle_auth_required(event, cdp_session)

        cdp_session.send.assert_called_once_with(
            "Fetch.continueWithAuth",
            {
                "requestId": "req-4",
                "authChallengeResponse": {"response": "CancelAuth"},
            },
        )

    @pytest.mark.asyncio
    async def test_auth_error_does_not_raise(self) -> None:
        """Errors during auth handling should be caught, not raised."""
        interceptor = self._make_interceptor(proxy_username="user1", proxy_password="pass1")
        cdp_session = self._make_cdp_session()
        cdp_session.send.side_effect = Exception("CDP connection lost")

        event = {
            "requestId": "req-5",
            "authChallenge": {"source": "Proxy", "origin": "http://proxy.example.com"},
            "request": {"url": "https://example.com/page"},
        }

        # Should not raise
        await interceptor._handle_auth_required(event, cdp_session)

    def test_init_stores_proxy_credentials(self) -> None:
        """Constructor should store proxy credentials."""
        interceptor = self._make_interceptor(proxy_username="user", proxy_password="pass")
        assert interceptor._proxy_username == "user"
        assert interceptor._proxy_password == "pass"

    def test_init_no_proxy_credentials(self) -> None:
        """Constructor without credentials should store None."""
        interceptor = self._make_interceptor()
        assert interceptor._proxy_username is None
        assert interceptor._proxy_password is None

    @pytest.mark.asyncio
    async def test_enable_for_page_with_proxy_auth(self) -> None:
        """enable_for_page with credentials should add Request-stage pattern and authRequired handler."""
        interceptor = self._make_interceptor(proxy_username="user", proxy_password="pass")

        mock_cdp_session = self._make_cdp_session()
        mock_page = MagicMock()
        mock_page.url = "about:blank"
        mock_page.context.new_cdp_session = AsyncMock(return_value=mock_cdp_session)

        await interceptor.enable_for_page(mock_page)

        # Verify Fetch.enable with both Response (downloads) and Request (auth) patterns
        mock_cdp_session.send.assert_called_once_with(
            "Fetch.enable",
            {
                "patterns": [
                    {"requestStage": "Response"},
                    {"urlPattern": "*", "requestStage": "Request"},
                ],
                "handleAuthRequests": True,
            },
        )

        # Verify both handlers registered
        event_names = [call.args[0] for call in mock_cdp_session.on.call_args_list]
        assert "Fetch.requestPaused" in event_names
        assert "Fetch.authRequired" in event_names

    @pytest.mark.asyncio
    async def test_enable_for_page_without_proxy_auth(self) -> None:
        """enable_for_page without credentials should only intercept Response stage."""
        interceptor = self._make_interceptor()

        mock_cdp_session = self._make_cdp_session()
        mock_page = MagicMock()
        mock_page.url = "about:blank"
        mock_page.context.new_cdp_session = AsyncMock(return_value=mock_cdp_session)

        await interceptor.enable_for_page(mock_page)

        # Verify Fetch.enable with Response-only pattern, no auth
        mock_cdp_session.send.assert_called_once_with(
            "Fetch.enable",
            {
                "patterns": [{"requestStage": "Response"}],
                "handleAuthRequests": False,
            },
        )

        # Verify only requestPaused handler (no authRequired)
        event_names = [call.args[0] for call in mock_cdp_session.on.call_args_list]
        assert "Fetch.requestPaused" in event_names
        assert "Fetch.authRequired" not in event_names

    @pytest.mark.asyncio
    async def test_request_stage_continues_request(self) -> None:
        """Request-stage events (no responseStatusCode) should be continued with Fetch.continueRequest."""
        interceptor = self._make_interceptor(proxy_username="user", proxy_password="pass")
        cdp_session = self._make_cdp_session()

        event = {
            "requestId": "req-1",
            "request": {"url": "https://example.com/page"},
            "resourceType": "Document",
            # No responseStatusCode — this is a Request-stage event
        }

        await interceptor._handle_request_paused(event, cdp_session)

        cdp_session.send.assert_called_once_with("Fetch.continueRequest", {"requestId": "req-1"})

    @pytest.mark.asyncio
    async def test_request_stage_error_does_not_retry(self) -> None:
        """Request-stage errors should not attempt recovery (no duplicate continueRequest)."""
        interceptor = self._make_interceptor(proxy_username="user", proxy_password="pass")
        cdp_session = self._make_cdp_session()

        cdp_session.send.side_effect = Exception("continueRequest failed")

        event = {
            "requestId": "req-err",
            "request": {"url": "https://example.com/page"},
            "resourceType": "Document",
            # No responseStatusCode — Request-stage event
        }

        await interceptor._handle_request_paused(event, cdp_session)

        # Only one call: the original continueRequest that failed. No recovery attempt.
        assert cdp_session.send.call_count == 1
        assert cdp_session.send.call_args.args[0] == "Fetch.continueRequest"

    @pytest.mark.asyncio
    async def test_malformed_event_missing_request_id(self) -> None:
        """Malformed event without requestId should be caught, not raise."""
        interceptor = self._make_interceptor(proxy_username="user1", proxy_password="pass1")
        cdp_session = self._make_cdp_session()

        event: dict = {
            "authChallenge": {"source": "Proxy", "origin": "http://proxy.example.com"},
            "request": {"url": "https://example.com/page"},
        }

        # Should not raise — KeyError is caught by the try/except
        await interceptor._handle_auth_required(event, cdp_session)
        cdp_session.send.assert_not_called()

    @pytest.mark.asyncio
    async def test_retry_loop_prevention(self) -> None:
        """Second auth attempt for the same requestId should CancelAuth to prevent infinite loop."""
        interceptor = self._make_interceptor(proxy_username="user1", proxy_password="pass1")
        cdp_session = self._make_cdp_session()

        event = {
            "requestId": "req-retry",
            "authChallenge": {"source": "Proxy", "origin": "http://proxy.example.com"},
            "request": {"url": "https://example.com/page"},
        }

        # First attempt: should provide credentials
        await interceptor._handle_auth_required(event, cdp_session)
        first_call = cdp_session.send.call_args
        assert first_call.args[1]["authChallengeResponse"]["response"] == "ProvideCredentials"

        cdp_session.send.reset_mock()

        # Second attempt (credentials rejected): should cancel
        await interceptor._handle_auth_required(event, cdp_session)
        second_call = cdp_session.send.call_args
        assert second_call.args[1]["authChallengeResponse"]["response"] == "CancelAuth"