crewAI/lib/crewai-tools/tests/utilities/test_safe_requests.py

"""Tests for redirect-aware safe HTTP helpers."""

from __future__ import annotations

import socket
from io import BytesIO
from typing import Any

import pytest
import requests

from crewai_tools.security.safe_requests import safe_get


def _response(url: str, status_code: int, *, location: str | None = None) -> requests.Response:
    response = requests.Response()
    response.status_code = status_code
    response.url = url
    response._content = b"ok"
    response.raw = BytesIO()
    if location is not None:
        response.headers["Location"] = location
    return response


@pytest.fixture
def public_dns(monkeypatch: pytest.MonkeyPatch) -> None:
    original_getaddrinfo = socket.getaddrinfo

    def fake_getaddrinfo(
        host: str, port: int, *args: Any, **kwargs: Any
    ) -> list[tuple[Any, ...]]:
        if host in {"public.example", "safe.example"}:
            return [
                (
                    socket.AF_INET,
                    socket.SOCK_STREAM,
                    6,
                    "",
                    ("93.184.216.34", port),
                )
            ]
        return original_getaddrinfo(host, port, *args, **kwargs)

    monkeypatch.setattr(socket, "getaddrinfo", fake_getaddrinfo)


def test_safe_get_blocks_direct_internal_url() -> None:
    with pytest.raises(ValueError, match="private/reserved IP"):
        safe_get("http://127.0.0.1/admin", timeout=15)


def _mock_get(monkeypatch: pytest.MonkeyPatch, get_response: Any) -> None:
    monkeypatch.setattr(
        "crewai_tools.security.safe_requests.requests.get",
        get_response,
    )


def test_safe_get_blocks_redirect_to_internal_url(
    monkeypatch: pytest.MonkeyPatch, public_dns: None
) -> None:
    requested_urls: list[str] = []

    def fake_get(url: str, **kwargs: Any) -> requests.Response:
        requested_urls.append(url)
        assert kwargs["allow_redirects"] is False
        return _response(url, 302, location="http://127.0.0.1/admin")

    _mock_get(monkeypatch, fake_get)

    with pytest.raises(ValueError, match="private/reserved IP"):
        safe_get("http://public.example/start", timeout=15)

    assert requested_urls == ["http://public.example/start"]


def test_safe_get_follows_safe_relative_redirect(
    monkeypatch: pytest.MonkeyPatch, public_dns: None
) -> None:
    requested_urls: list[str] = []

    def fake_get(url: str, **kwargs: Any) -> requests.Response:
        requested_urls.append(url)
        assert kwargs["allow_redirects"] is False
        if url == "http://public.example/start":
            return _response(url, 302, location="/final")
        return _response(url, 200)

    _mock_get(monkeypatch, fake_get)

    response = safe_get("http://public.example/start", timeout=15)

    assert response.status_code == 200
    assert response.url == "http://public.example/final"
    assert requested_urls == [
        "http://public.example/start",
        "http://public.example/final",
    ]
    assert len(response.history) == 1


def test_safe_get_fails_closed_after_too_many_redirects(
    monkeypatch: pytest.MonkeyPatch, public_dns: None
) -> None:
    def fake_get(url: str, **kwargs: Any) -> requests.Response:
        return _response(url, 302, location="http://safe.example/again")

    _mock_get(monkeypatch, fake_get)

    with pytest.raises(ValueError, match="Too many redirects"):
        safe_get("http://public.example/start", max_redirects=1, timeout=15)


def _closable_response(
    url: str, status_code: int, *, location: str | None = None, closed: list[str]
) -> requests.Response:
    """Build a response that records its own URL when closed."""
    response = _response(url, status_code, location=location)
    response.close = lambda: closed.append(url)  # type: ignore[method-assign]
    return response


def test_safe_get_closes_earlier_hops_after_too_many_redirects(
    monkeypatch: pytest.MonkeyPatch, public_dns: None
) -> None:
    """Hops accumulated before the failure must not be left open.

    Under stream=True each hop holds its connection until its body is read or
    closed, and a caller handed an exception has no handle on them.
    """
    closed: list[str] = []

    def fake_get(url: str, **kwargs: Any) -> requests.Response:
        return _closable_response(
            url, 302, location="http://safe.example/again", closed=closed
        )

    _mock_get(monkeypatch, fake_get)

    with pytest.raises(ValueError, match="Too many redirects"):
        safe_get("http://public.example/start", max_redirects=2, timeout=15, stream=True)

    assert len(closed) == 3


def test_safe_get_closes_earlier_hops_when_a_redirect_is_rejected(
    monkeypatch: pytest.MonkeyPatch, public_dns: None
) -> None:
    """A hop rejected mid-chain still releases the connections already open."""
    closed: list[str] = []

    def fake_get(url: str, **kwargs: Any) -> requests.Response:
        if url == "http://public.example/start":
            return _closable_response(
                url, 302, location="http://safe.example/next", closed=closed
            )
        return _closable_response(
            url, 302, location="http://169.254.169.254/latest", closed=closed
        )

    _mock_get(monkeypatch, fake_get)

    with pytest.raises(ValueError, match="private/reserved IP"):
        safe_get("http://public.example/start", timeout=15, stream=True)

    assert closed == ["http://safe.example/next", "http://public.example/start"]


def test_safe_get_leaves_hops_open_on_success(
    monkeypatch: pytest.MonkeyPatch, public_dns: None
) -> None:
    """On success the hops belong to the caller, via response.history."""
    closed: list[str] = []

    def fake_get(url: str, **kwargs: Any) -> requests.Response:
        if url == "http://public.example/start":
            return _closable_response(url, 302, location="/final", closed=closed)
        return _closable_response(url, 200, closed=closed)

    _mock_get(monkeypatch, fake_get)

    response = safe_get("http://public.example/start", timeout=15, stream=True)

    assert closed == []
    assert len(response.history) == 1


def test_safe_get_strips_credentials_on_cross_origin_redirect(
    monkeypatch: pytest.MonkeyPatch, public_dns: None
) -> None:
    requests_made: list[tuple[str, dict[str, Any]]] = []

    def fake_get(url: str, **kwargs: Any) -> requests.Response:
        requests_made.append((url, kwargs))
        if url == "http://public.example/start":
            return _response(url, 302, location="http://safe.example/final")
        return _response(url, 200)

    _mock_get(monkeypatch, fake_get)

    response = safe_get(
        "http://public.example/start",
        timeout=15,
        headers={
            "Authorization": "Bearer token",
            "Authorization-Custom": "secret token",
            "Cookie": "session=abc",
            "X-API-Key": "api key",
            "X-CrewAI-Token": "crewai token",
            "User-Agent": "crewai-test",
        },
        cookies={"session": "abc"},
    )

    assert response.status_code == 200
    assert requests_made[0][1]["headers"] == {
        "Authorization": "Bearer token",
        "Authorization-Custom": "secret token",
        "Cookie": "session=abc",
        "X-API-Key": "api key",
        "X-CrewAI-Token": "crewai token",
        "User-Agent": "crewai-test",
    }
    assert requests_made[0][1]["cookies"] == {"session": "abc"}
    assert requests_made[1][1]["headers"] == {"User-Agent": "crewai-test"}
    assert "cookies" not in requests_made[1][1]


def test_safe_get_preserves_credentials_on_same_origin_redirect(
    monkeypatch: pytest.MonkeyPatch, public_dns: None
) -> None:
    requests_made: list[tuple[str, dict[str, Any]]] = []

    def fake_get(url: str, **kwargs: Any) -> requests.Response:
        requests_made.append((url, kwargs))
        if url == "http://public.example/start":
            return _response(url, 302, location="/final")
        return _response(url, 200)

    _mock_get(monkeypatch, fake_get)

    safe_get(
        "http://public.example/start",
        timeout=15,
        headers={"Authorization": "Bearer token"},
        cookies={"session": "abc"},
    )

    assert requests_made[1][1]["headers"] == {"Authorization": "Bearer token"}
    assert requests_made[1][1]["cookies"] == {"session": "abc"}