Files
crewAI/lib/crewai-tools/tests/utilities/test_safe_requests.py
2026-06-25 11:27:25 +08:00

115 lines
3.4 KiB
Python

"""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 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")
monkeypatch.setattr(
"crewai_tools.security.safe_requests.requests.get",
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)
monkeypatch.setattr(
"crewai_tools.security.safe_requests.requests.get",
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")
monkeypatch.setattr(
"crewai_tools.security.safe_requests.requests.get",
fake_get,
)
with pytest.raises(ValueError, match="Too many redirects"):
safe_get("http://public.example/start", max_redirects=1, timeout=15)