mirror of
https://github.com/pewdiepie-archdaemon/odysseus.git
synced 2026-06-16 17:55:26 -04:00
RaresKeY
105 lines
2.9 KiB
Python
105 lines
2.9 KiB
Python
from src import ai_interaction
|
|
|
|
|
|
class _GenerationResponse:
|
|
status_code = 200
|
|
text = ""
|
|
|
|
def __init__(self, image_url):
|
|
self._image_url = image_url
|
|
|
|
def json(self):
|
|
return {"data": [{"url": self._image_url}]}
|
|
|
|
|
|
class _DownloadResponse:
|
|
status_code = 503
|
|
content = b""
|
|
|
|
|
|
def _patch_generation(monkeypatch, image_url):
|
|
async def _post(self, url, json, headers):
|
|
return _GenerationResponse(image_url)
|
|
|
|
class _AsyncClient:
|
|
def __init__(self, *args, **kwargs):
|
|
pass
|
|
|
|
async def __aenter__(self):
|
|
return self
|
|
|
|
async def __aexit__(self, *exc):
|
|
return False
|
|
|
|
post = _post
|
|
|
|
import httpx
|
|
import src.settings as settings
|
|
|
|
monkeypatch.setattr(settings, "load_settings", lambda: {})
|
|
monkeypatch.setattr(httpx, "AsyncClient", _AsyncClient)
|
|
monkeypatch.setattr(
|
|
ai_interaction,
|
|
"_resolve_model",
|
|
lambda model_spec, owner=None: (
|
|
"https://api.openai.example/v1/chat/completions",
|
|
"dall-e-3",
|
|
{"Authorization": "Bearer test"},
|
|
),
|
|
)
|
|
|
|
|
|
async def test_generate_image_validates_provider_url_before_download(monkeypatch):
|
|
import httpx
|
|
import src.url_safety as url_safety
|
|
|
|
provider_url = "https://images.example.com/generated.png?sig=abc"
|
|
events = []
|
|
_patch_generation(monkeypatch, provider_url)
|
|
|
|
def _check_outbound_url(url, *, block_private=False):
|
|
events.append(("check", url, block_private))
|
|
return True, "ok"
|
|
|
|
def _get(url, *, timeout):
|
|
events.append(("get", url, timeout))
|
|
return _DownloadResponse()
|
|
|
|
monkeypatch.setattr(url_safety, "check_outbound_url", _check_outbound_url)
|
|
monkeypatch.setattr(httpx, "get", _get)
|
|
|
|
result = await ai_interaction.do_generate_image("draw a chair\ndall-e-3")
|
|
|
|
assert result["image_url"] == provider_url
|
|
assert events == [
|
|
("check", provider_url, False),
|
|
("get", provider_url, 60),
|
|
]
|
|
|
|
|
|
async def test_generate_image_rejects_unsafe_provider_url_without_download(monkeypatch):
|
|
import httpx
|
|
import src.url_safety as url_safety
|
|
|
|
unsafe_url = "http://169.254.169.254/latest/meta-data"
|
|
events = []
|
|
_patch_generation(monkeypatch, unsafe_url)
|
|
|
|
def _check_outbound_url(url, *, block_private=False):
|
|
events.append(("check", url, block_private))
|
|
return False, "link-local address blocked (SSRF metadata risk): 169.254.169.254"
|
|
|
|
def _get(url, *, timeout):
|
|
raise AssertionError("unsafe provider image URL must not be downloaded")
|
|
|
|
monkeypatch.setattr(url_safety, "check_outbound_url", _check_outbound_url)
|
|
monkeypatch.setattr(httpx, "get", _get)
|
|
|
|
result = await ai_interaction.do_generate_image("draw a chair\ndall-e-3")
|
|
|
|
assert result["error"] == (
|
|
"Image API returned unsafe image URL: "
|
|
"link-local address blocked (SSRF metadata risk): 169.254.169.254"
|
|
)
|
|
assert events == [("check", unsafe_url, False)]
|