Oleksandr Kozachuk
379 lines
17 KiB
Python
379 lines
17 KiB
Python
"""Unit coverage for SPEC-011 URL reading (URL-01..09)."""
|
|
|
|
import json
|
|
import unittest
|
|
from unittest.mock import AsyncMock, Mock, patch
|
|
|
|
from fjerkroa_bot.openai_responder import OpenAIResponder
|
|
from fjerkroa_bot.url_reader import FETCH_URL_TOOL, URLReader, guard_url
|
|
|
|
from .test_bdd_envelope import envelope
|
|
|
|
CONFIG = {"openai-token": "t", "model": "m", "system": "s", "history-limit": 5}
|
|
|
|
|
|
class TestToolOffered(unittest.TestCase):
|
|
def test_tool_present_only_when_enabled(self):
|
|
"""URL-01: fetch_url appears in the tool list only with enable-url-reading."""
|
|
off = OpenAIResponder(CONFIG, "chat")
|
|
self.assertNotIn("fetch_url", [f["name"] for f in off._available_tools()])
|
|
on = OpenAIResponder(dict(CONFIG, **{"enable-url-reading": True}), "chat")
|
|
self.assertIn("fetch_url", [f["name"] for f in on._available_tools()])
|
|
self.assertEqual(FETCH_URL_TOOL["name"], "fetch_url")
|
|
|
|
|
|
class TestSchemeGuard(unittest.TestCase):
|
|
def test_non_http_schemes_refused(self):
|
|
"""URL-02: only http/https pass the guard."""
|
|
self.assertIsNone(guard_url("https://example.com/article"))
|
|
for bad in ("file:///etc/passwd", "ftp://host/x", "data:text/html,x", "gopher://h", "no-scheme.com/x"):
|
|
self.assertIsNotNone(guard_url(bad))
|
|
|
|
|
|
class TestSSRFGuard(unittest.TestCase):
|
|
def test_private_and_loopback_refused(self):
|
|
"""URL-03: private/loopback/link-local literals are refused without DNS."""
|
|
for bad in (
|
|
"http://127.0.0.1/admin",
|
|
"http://localhost/x", # resolves to loopback
|
|
"http://10.0.0.5/x",
|
|
"http://192.168.1.1/x",
|
|
"http://169.254.169.254/latest/meta-data", # cloud metadata
|
|
"http://[::1]/x",
|
|
):
|
|
self.assertIsNotNone(guard_url(bad), f"{bad} should be refused")
|
|
|
|
def test_public_ip_allowed(self):
|
|
"""URL-03: a public IP literal passes."""
|
|
self.assertIsNone(guard_url("http://93.184.216.34/"))
|
|
|
|
@patch("fjerkroa_bot.url_reader.socket.getaddrinfo")
|
|
def test_dns_to_private_refused(self, getaddrinfo):
|
|
"""URL-03: a hostname resolving to a private IP is refused."""
|
|
getaddrinfo.return_value = [(2, 1, 6, "", ("10.1.2.3", 0))]
|
|
self.assertIsNotNone(guard_url("http://evil.example.com/x"))
|
|
|
|
|
|
class TestRedirectRevalidation(unittest.IsolatedAsyncioTestCase):
|
|
async def test_redirect_to_internal_refused(self):
|
|
"""URL-04: a public URL redirecting to localhost is refused at the hop."""
|
|
reader = URLReader(lambda: {}, None)
|
|
|
|
class FakeResp:
|
|
status = 302
|
|
headers = {"Location": "http://127.0.0.1/secret"}
|
|
url = "http://safe.example.com"
|
|
|
|
async def __aenter__(self):
|
|
return self
|
|
|
|
async def __aexit__(self, *a):
|
|
return False
|
|
|
|
def raise_for_status(self):
|
|
pass
|
|
|
|
class FakeSession:
|
|
def get(self, url, allow_redirects=False):
|
|
return FakeResp()
|
|
|
|
with patch("fjerkroa_bot.url_reader.guard_url", side_effect=[None, "refused internal"]):
|
|
with self.assertRaises(ValueError):
|
|
await reader._get(FakeSession(), "http://safe.example.com", 1000)
|
|
|
|
|
|
class TestMetaRefresh(unittest.IsolatedAsyncioTestCase):
|
|
async def test_follows_meta_refresh_to_real_article(self):
|
|
"""URL-04: a getnews-style meta-refresh stub is followed to the real article."""
|
|
reader = URLReader(lambda: {}, None)
|
|
stub = (
|
|
b'<html><head><meta http-equiv="refresh" content="0;url=https://pushsquare.com/real"></head><body>Redirecting...</body></html>'
|
|
)
|
|
article = b"<html><body><h1>MARVEL Tokon</h1><p>Full article text here</p></body></html>"
|
|
calls = []
|
|
|
|
async def fake_get(session, url, max_bytes):
|
|
calls.append(url)
|
|
return (url, stub if "stub" in url else article, "text/html")
|
|
|
|
reader._get = fake_get # type: ignore
|
|
with patch("fjerkroa_bot.url_reader.guard_url", return_value=None):
|
|
import fjerkroa_bot.url_reader as ur
|
|
|
|
# patch the session context so fetch() runs against fake_get
|
|
class FakeCM:
|
|
async def __aenter__(self):
|
|
return object()
|
|
|
|
async def __aexit__(self, *a):
|
|
return False
|
|
|
|
with patch.object(ur.aiohttp, "ClientSession", return_value=FakeCM()):
|
|
result = await reader.fetch("https://gggemein.de/url/stub.html", "chat", "alice")
|
|
self.assertIn("Full article text", result["text"])
|
|
self.assertEqual(result["url"], "https://pushsquare.com/real")
|
|
self.assertIn("https://pushsquare.com/real", calls)
|
|
|
|
async def test_meta_refresh_to_internal_is_not_followed(self):
|
|
"""URL-04: a meta-refresh pointing at an internal IP is refused (SSRF)."""
|
|
reader = URLReader(lambda: {}, None)
|
|
stub = b'<meta http-equiv="refresh" content="0; url=http://127.0.0.1/secret">Redirecting'
|
|
|
|
async def fake_get(session, url, max_bytes):
|
|
return (url, stub, "text/html")
|
|
|
|
reader._get = fake_get # type: ignore
|
|
import fjerkroa_bot.url_reader as ur
|
|
|
|
class FakeCM:
|
|
async def __aenter__(self):
|
|
return object()
|
|
|
|
async def __aexit__(self, *a):
|
|
return False
|
|
|
|
def guard(u):
|
|
return "refused" if "127.0.0.1" in u else None
|
|
|
|
with patch("fjerkroa_bot.url_reader.guard_url", side_effect=guard):
|
|
with patch.object(ur.aiohttp, "ClientSession", return_value=FakeCM()):
|
|
result = await reader.fetch("https://safe.com/x", "chat", "alice")
|
|
self.assertEqual(result["url"], "https://safe.com/x") # did not follow to 127.0.0.1
|
|
|
|
|
|
class TestTextExtraction(unittest.TestCase):
|
|
def test_html_reduced_to_text(self):
|
|
"""URL-05: scripts/styles dropped, tags stripped."""
|
|
reader = URLReader(lambda: {}, None)
|
|
html = "<html><head><style>x{}</style></head><body><h1>Titel</h1><script>evil()</script><p>Inhalt hier</p></body></html>"
|
|
text = reader._to_text(html)
|
|
self.assertIn("Titel", text)
|
|
self.assertIn("Inhalt hier", text)
|
|
self.assertNotIn("evil", text)
|
|
self.assertNotIn("x{}", text)
|
|
|
|
def test_chrome_and_link_boilerplate_dropped(self):
|
|
"""URL-08: nav/header/footer skipped; short link-dominated blocks (menus, related lists) removed."""
|
|
reader = URLReader(lambda: {}, None)
|
|
html = (
|
|
"<html><body>"
|
|
"<nav><a href='/a'>Home</a> <a href='/b'>Games</a></nav>"
|
|
"<header><a href='/login'>Login</a></header>"
|
|
"<ul><li><a href='/1'>Related article one</a></li><li><a href='/2'>Related article two</a></li></ul>"
|
|
"<article><p>The pop-up event runs from August 4 in Shibuya, with details "
|
|
"<a href='/x'>on the official page</a> for anyone attending the exhibition.</p></article>"
|
|
"<footer><a href='/imprint'>Imprint</a></footer>"
|
|
"</body></html>"
|
|
)
|
|
text = reader._to_text(html)
|
|
self.assertIn("pop-up event", text)
|
|
self.assertIn("on the official page", text) # inline link in a real paragraph survives
|
|
for chrome in ("Home", "Login", "Related article one", "Imprint"):
|
|
self.assertNotIn(chrome, text)
|
|
|
|
def test_default_cap_is_8000(self):
|
|
"""URL-08: the default url-max-chars budget is 8000."""
|
|
from fjerkroa_bot.url_reader import DEFAULT_MAX_CHARS
|
|
|
|
self.assertEqual(DEFAULT_MAX_CHARS, 8000)
|
|
|
|
|
|
class TestBodyReadCollectsAllChunks(unittest.IsolatedAsyncioTestCase):
|
|
async def test_get_reads_past_first_chunk(self):
|
|
"""URL-05 regression: body arrives in many chunks; all are collected up to the cap."""
|
|
reader = URLReader(lambda: {}, None)
|
|
chunks = [b"<title>t</title>", b"<p>middle</p>", b"<p>end</p>"]
|
|
|
|
class FakeContent:
|
|
@staticmethod
|
|
async def iter_chunked(size):
|
|
for chunk in chunks:
|
|
yield chunk
|
|
|
|
class FakeResp:
|
|
status = 200
|
|
headers = {}
|
|
url = "http://safe.example.com"
|
|
content = FakeContent()
|
|
|
|
async def __aenter__(self):
|
|
return self
|
|
|
|
async def __aexit__(self, *a):
|
|
return False
|
|
|
|
def raise_for_status(self):
|
|
pass
|
|
|
|
class FakeSession:
|
|
def get(self, url, allow_redirects=False):
|
|
return FakeResp()
|
|
|
|
with patch("fjerkroa_bot.url_reader.guard_url", return_value=None):
|
|
_, body, _ = await reader._get(FakeSession(), "http://safe.example.com", 1000)
|
|
self.assertEqual(body, b"".join(chunks))
|
|
|
|
with patch("fjerkroa_bot.url_reader.guard_url", return_value=None):
|
|
_, body, _ = await reader._get(FakeSession(), "http://safe.example.com", 20)
|
|
self.assertEqual(body, b"".join(chunks)[:20])
|
|
|
|
|
|
class TestFetchSanitizes(unittest.IsolatedAsyncioTestCase):
|
|
async def test_fetch_result_is_sanitized_and_capped(self):
|
|
"""URL-05: fetch output is length-capped and @everyone-neutralized."""
|
|
reader = URLReader(lambda: {"url-max-chars": 50}, None)
|
|
payload = ("<p>@everyone " + "x" * 5000 + "</p>").encode()
|
|
with patch.object(reader, "_get", new=AsyncMock(return_value=("http://x.com", payload, "text/html"))):
|
|
result = await reader.fetch("http://x.com", "chat", "alice")
|
|
self.assertLessEqual(len(result["text"]), 50)
|
|
self.assertNotIn("@everyone", result["text"])
|
|
|
|
async def test_fetch_error_is_reported_not_raised(self):
|
|
"""URL-05: a fetch failure returns an error dict the model can relay."""
|
|
reader = URLReader(lambda: {}, None)
|
|
with patch.object(reader, "_get", new=AsyncMock(side_effect=ValueError("refused non-public address"))):
|
|
result = await reader.fetch("http://10.0.0.1", "chat", "alice")
|
|
self.assertIn("error", result)
|
|
|
|
|
|
class TestImageIngest(unittest.IsolatedAsyncioTestCase):
|
|
async def test_page_images_go_to_cache_ssrf_guarded(self):
|
|
"""URL-06: og:image + <img> ingested (cap honored), internal srcs skipped."""
|
|
cache = type("C", (), {})()
|
|
cache.ingest_url = AsyncMock(side_effect=["sha1", "sha2", "sha3"])
|
|
cache.recent = Mock(return_value=[{"sha256": "sha1", "ext": "jpg"}, {"sha256": "sha2", "ext": "png"}])
|
|
cache.data_url = Mock(side_effect=lambda sha, ext: f"data:image/{ext};base64,{sha}")
|
|
reader = URLReader(lambda: {"url-max-images": 2}, cache)
|
|
html = (
|
|
'<meta property="og:image" content="https://cdn.example.com/hero.jpg">'
|
|
'<img src="https://cdn.example.com/a.png"><img src="http://127.0.0.1/internal.png">'
|
|
)
|
|
|
|
# guard by scheme/loopback only, no real DNS in the test
|
|
def fake_guard(url):
|
|
return "refused" if "127.0.0.1" in url else None
|
|
|
|
with patch("fjerkroa_bot.url_reader.guard_url", side_effect=fake_guard):
|
|
data_urls = await reader._ingest_images(html, "https://example.com", "chat", "alice")
|
|
self.assertEqual(len(data_urls), 2) # og:image + first public img, cap 2
|
|
self.assertEqual(data_urls[0], "data:image/jpg;base64,sha1") # URL-09: data URLs for vision
|
|
ingested = [call.args[0] for call in cache.ingest_url.await_args_list]
|
|
self.assertNotIn("http://127.0.0.1/internal.png", ingested)
|
|
|
|
|
|
class TestPerUserCap(unittest.IsolatedAsyncioTestCase):
|
|
async def test_dispatch_caps_fetches(self):
|
|
"""URL-07: over url-daily-per-user, fetch_url returns an error without fetching."""
|
|
responder = OpenAIResponder(dict(CONFIG, **{"enable-url-reading": True, "url-daily-per-user": 2}), "chat")
|
|
responder.url_reader.fetch = AsyncMock(return_value={"url": "x", "text": "ok"})
|
|
for _ in range(2):
|
|
await responder._dispatch_tool("fetch_url", {"url": "http://x.com"}, "alice")
|
|
blocked = await responder._dispatch_tool("fetch_url", {"url": "http://x.com"}, "alice")
|
|
self.assertIn("error", blocked)
|
|
self.assertEqual(responder.url_reader.fetch.await_count, 2)
|
|
|
|
|
|
class TestFetchedImagesBecomeVision(unittest.IsolatedAsyncioTestCase):
|
|
"""URL-09: fetch results carry vision data URLs, direct image URLs are ingested."""
|
|
|
|
@staticmethod
|
|
def _cache():
|
|
cache = Mock()
|
|
cache.ingest_url = AsyncMock(return_value="abc123")
|
|
cache.ingest_bytes = Mock(return_value="abc123")
|
|
cache.recent = Mock(return_value=[{"sha256": "abc123", "ext": "png"}])
|
|
cache.data_url = Mock(return_value="data:image/png;base64,AAA")
|
|
return cache
|
|
|
|
@staticmethod
|
|
def _session_cm():
|
|
import fjerkroa_bot.url_reader as ur
|
|
|
|
class FakeCM:
|
|
async def __aenter__(self):
|
|
return object()
|
|
|
|
async def __aexit__(self, *a):
|
|
return False
|
|
|
|
return patch.object(ur.aiohttp, "ClientSession", return_value=FakeCM())
|
|
|
|
async def test_html_page_vision_data_urls(self):
|
|
"""URL-09: og:image lands in the result's vision list, count matches."""
|
|
reader = URLReader(lambda: {}, self._cache())
|
|
html = b'<meta property="og:image" content="https://x.com/c.png"><p>Comic of the day, longer text.</p>'
|
|
|
|
async def fake_get(session, url, max_bytes):
|
|
return (url, html, "text/html")
|
|
|
|
reader._get = fake_get # type: ignore
|
|
with patch("fjerkroa_bot.url_reader.guard_url", return_value=None):
|
|
with self._session_cm():
|
|
result = await reader.fetch("https://xkcd.com/1234", "chat", "alice")
|
|
self.assertEqual(result["vision"], ["data:image/png;base64,AAA"])
|
|
self.assertEqual(result["images_cached"], 1)
|
|
|
|
async def test_direct_image_url_ingested(self):
|
|
"""URL-09: content-type image/* -> direct ingest, text '(image)'."""
|
|
cache = self._cache()
|
|
reader = URLReader(lambda: {}, cache)
|
|
|
|
async def fake_get(session, url, max_bytes):
|
|
return (url, b"\x89PNG-bytes", "image/png")
|
|
|
|
reader._get = fake_get # type: ignore
|
|
with patch("fjerkroa_bot.url_reader.guard_url", return_value=None):
|
|
with self._session_cm():
|
|
result = await reader.fetch("https://imgs.xkcd.com/comics/x.png", "chat", "alice")
|
|
self.assertEqual(result["text"], "(image)")
|
|
self.assertEqual(result["vision"], ["data:image/png;base64,AAA"])
|
|
cache.ingest_bytes.assert_called_once()
|
|
|
|
async def test_capped_image_body_not_ingested(self):
|
|
"""URL-09: an image body at the byte cap may be truncated - not ingested."""
|
|
cache = self._cache()
|
|
reader = URLReader(lambda: {"url-max-bytes": 10}, cache)
|
|
|
|
async def fake_get(session, url, max_bytes):
|
|
return (url, b"0123456789", "image/png") # len == cap
|
|
|
|
reader._get = fake_get # type: ignore
|
|
with patch("fjerkroa_bot.url_reader.guard_url", return_value=None):
|
|
with self._session_cm():
|
|
result = await reader.fetch("https://x.com/big.png", "chat", "alice")
|
|
self.assertEqual(result["vision"], [])
|
|
cache.ingest_bytes.assert_not_called()
|
|
|
|
|
|
class TestLegacyPathVision(unittest.IsolatedAsyncioTestCase):
|
|
async def test_legacy_tool_loop_appends_image_message(self):
|
|
"""URL-09: legacy path - vision data URLs become an image_url user message; never JSON text."""
|
|
responder = OpenAIResponder(dict(CONFIG, **{"enable-url-reading": True}), "chat")
|
|
responder._dispatch_tool = AsyncMock(
|
|
return_value={"url": "u", "text": "t", "images_cached": 1, "vision": ["data:image/png;base64,AAA"]}
|
|
)
|
|
func = Mock()
|
|
func.name = "fetch_url"
|
|
func.arguments = json.dumps({"url": "https://xkcd.com/1"})
|
|
call = Mock(id="tc1", type="function", function=func)
|
|
first_msg = Mock(content=None, role="assistant", tool_calls=[call], refusal=None)
|
|
first = Mock(choices=[Mock(message=first_msg)], usage=None)
|
|
final_msg = Mock(content=envelope(answer="seen", answer_needed=True), role="assistant", tool_calls=None, refusal=None)
|
|
final = Mock(choices=[Mock(message=final_msg)], usage=None)
|
|
with patch("fjerkroa_bot.openai_responder.openai_chat", new_callable=AsyncMock) as chat_mock:
|
|
chat_mock.side_effect = [first, final]
|
|
answer, _ = await responder.chat([{"role": "user", "content": "look at this"}], 10)
|
|
self.assertEqual(json.loads(answer["content"])["answer"], "seen")
|
|
final_messages = chat_mock.await_args_list[1].kwargs["messages"]
|
|
image_parts = [
|
|
part
|
|
for msg in final_messages
|
|
if isinstance(msg.get("content"), list)
|
|
for part in msg["content"]
|
|
if part.get("type") == "image_url"
|
|
]
|
|
self.assertEqual(image_parts[0]["image_url"]["url"], "data:image/png;base64,AAA")
|
|
tool_texts = [msg["content"] for msg in final_messages if msg.get("role") == "tool"]
|
|
self.assertNotIn("data:image", tool_texts[0]) # data URL never in JSON tool text
|
|
self.assertNotIn("vision", tool_texts[0])
|