File size: 9,309 Bytes
187966e b81fe79 cd1d8fb 4778987 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 | from __future__ import annotations
from pathlib import Path
from app.rag.extract import extract_file, extract_html
def _make_txt(d: Path) -> Path:
p = d / "doc.txt"
p.write_text("El producto pesa 25 kg y tiene garantía de 2 años.", encoding="utf-8")
return p
def _make_docx(d: Path) -> Path:
from docx import Document
doc = Document()
doc.add_paragraph("La garantía es de dos años.")
doc.add_paragraph("Envíos gratis a partir de 199€.")
p = d / "doc.docx"
doc.save(p)
return p
def _make_xlsx(d: Path) -> Path:
from openpyxl import Workbook
wb = Workbook()
ws = wb.active
ws.append(["Modelo", "X100"])
ws.append(["Peso", "25 kg"])
p = d / "doc.xlsx"
wb.save(p)
return p
def _make_pdf(d: Path) -> Path:
from fpdf import FPDF
pdf = FPDF()
pdf.add_page()
pdf.set_font("helvetica", size=12)
pdf.cell(0, 10, text="Manual de usuario TOORX")
p = d / "doc.pdf"
pdf.output(str(p))
return p
def test_extract_txt(tmp_path):
assert "25 kg" in extract_file(_make_txt(tmp_path))
def test_extract_docx(tmp_path):
text = extract_file(_make_docx(tmp_path))
assert "garantía es de dos años" in text
assert "199" in text
def test_extract_xlsx(tmp_path):
text = extract_file(_make_xlsx(tmp_path))
assert "X100" in text
assert "25 kg" in text
def test_extract_pdf(tmp_path):
text = extract_file(_make_pdf(tmp_path))
assert "Manual" in text
assert "TOORX" in text
def test_extract_html_strips_tags_and_scripts():
html = (
"<html><body><h1>Envíos</h1><p>Gratis a partir de 199€</p>"
"<script>evil()</script></body></html>"
)
text = extract_html(html)
assert "Envíos" in text
assert "Gratis a partir de 199€" in text
assert "evil" not in text
def test_extract_unsupported_raises(tmp_path):
p = tmp_path / "thing.xyz"
p.write_text("hi", encoding="utf-8")
import pytest
with pytest.raises(ValueError):
extract_file(p)
class _Resp:
def __init__(self, *, content=b"", text="", status_code=200, content_type="text/html"):
self.content = content
self.text = text
self.status_code = status_code
self.headers = {"content-type": content_type}
def raise_for_status(self):
if self.status_code >= 400:
raise RuntimeError(f"http {self.status_code}")
class _FakeClient:
"""Maps URLs to responses; unknown URLs return 404."""
def __init__(self, pages):
self.pages = pages
self.requested = []
async def get(self, url, headers=None):
self.requested.append(url)
return self.pages.get(url, _Resp(status_code=404))
async def test_crawl_follows_internal_links_and_sitemap():
from app.rag.extract import crawl_url
base = "https://shop.example"
pages = {
f"{base}/": _Resp(content=(
f'<html><body><h1>Inicio</h1>'
f'<a href="/faq">FAQ</a> <a href="{base}/envios">Envios</a>'
f'<a href="https://otro.com/x">externo</a></body></html>'
).encode()),
f"{base}/faq": _Resp(content=b"<html><body>Garantia de dos anios</body></html>"),
f"{base}/envios": _Resp(content=b"<html><body>Envio gratis desde 199</body></html>"),
# discovered only via sitemap, not linked from any page
f"{base}/oculta": _Resp(content=b"<html><body>Pagina secreta del sitemap</body></html>"),
f"{base}/sitemap.xml": _Resp(
text=f"<urlset><url><loc>{base}/oculta</loc></url></urlset>",
content_type="application/xml",
),
}
client = _FakeClient(pages)
text = await crawl_url(base + "/", client=client)
# every internal page is indexed (landing + linked + sitemap-only)
assert "Garantia de dos anios" in text
assert "Envio gratis desde 199" in text
assert "Pagina secreta del sitemap" in text
# external domain is never fetched
assert "https://otro.com/x" not in client.requested
async def test_crawl_respects_max_pages():
from app.rag.extract import crawl_url
base = "https://big.example"
pages = {f"{base}/sitemap.xml": _Resp(status_code=404)}
# landing links to 10 pages; cap at 3
links = " ".join(f'<a href="/p{i}">p{i}</a>' for i in range(10))
pages[f"{base}/"] = _Resp(content=f"<html><body>{links}</body></html>".encode())
for i in range(10):
pages[f"{base}/p{i}"] = _Resp(content=f"<html><body>page {i}</body></html>".encode())
client = _FakeClient(pages)
await crawl_url(base + "/", max_pages=3, client=client)
# sitemap probe + at most 3 page fetches
page_fetches = [u for u in client.requested if "sitemap" not in u]
assert len(page_fetches) <= 3
# --- crawl-time SSRF guard (for untrusted start URLs, e.g. the public demo) --
def test_assert_public_host_blocks_private_and_allows_public(monkeypatch):
import pytest
from app.rag import extract
# IP literals are classified directly, no DNS — incl. IPv4-mapped IPv6.
for bad in ("127.0.0.1", "10.0.0.5", "169.254.169.254", "::1", "::ffff:127.0.0.1", ""):
with pytest.raises(extract._BlockedAddress):
extract._assert_public_host(bad)
# A hostname whose resolution includes a private answer is refused...
monkeypatch.setattr(extract, "_resolve_addrs", lambda h: ["93.184.216.34", "10.1.2.3"])
with pytest.raises(extract._BlockedAddress):
extract._assert_public_host("rebind.example")
# ...while an all-public hostname passes (no exception).
monkeypatch.setattr(extract, "_resolve_addrs", lambda h: ["93.184.216.34"])
extract._assert_public_host("good.example")
def _redirect_to_private_transport():
import httpx
def handler(request):
if request.url.host == "shop.test":
return httpx.Response(302, headers={"location": "http://169.254.169.254/secret"})
return httpx.Response(
200, text="CLOUD METADATA TOKEN abc123", headers={"content-type": "text/html"}
)
return httpx.MockTransport(handler)
async def test_crawl_guard_blocks_redirect_into_private_network(monkeypatch):
"""A public start URL that 302-redirects to the cloud-metadata endpoint must
NOT be followed: the internal content never enters the crawl corpus."""
import httpx
from app.rag import extract
monkeypatch.setattr(extract, "_resolve_addrs", lambda h: ["93.184.216.34"])
client = httpx.AsyncClient(transport=_redirect_to_private_transport(), follow_redirects=True)
text, pages = await extract.crawl_site("https://shop.test/", client=client, guard_ssrf=True)
await client.aclose()
assert "CLOUD METADATA TOKEN" not in text
assert pages == [] # the redirect was refused; nothing was fetched
async def test_crawl_without_guard_would_leak_redirect_target(monkeypatch):
"""Control proving the guard above is load-bearing: the SAME redirect WITHOUT
guard_ssrf leaks the internal content (so the block test is not vacuous)."""
import httpx
from app.rag import extract
client = httpx.AsyncClient(transport=_redirect_to_private_transport(), follow_redirects=True)
text, _pages = await extract.crawl_site("https://shop.test/", client=client)
await client.aclose()
assert "CLOUD METADATA TOKEN" in text
def _blocked_store_with_reader_transport():
"""Store returns 200 but EMPTY (datacenter IP soft-blocked); the r.jina.ai
reader returns the real rendered text from its own IPs."""
import httpx
def handler(request):
if request.url.host == "r.jina.ai":
return httpx.Response(
200,
text="Title: Tienda Frikadas\n\nVendemos juegos de mesa. Envios en 48h. "
"Devoluciones en 14 dias.",
headers={"content-type": "text/plain"},
)
# the store itself: reachable but serves nothing useful to the crawler
return httpx.Response(200, text="", headers={"content-type": "text/html"})
return httpx.MockTransport(handler)
async def test_remote_fallback_recovers_blocked_store():
"""When a direct crawl comes back empty and remote_fallback=True, crawl_site
recovers the store's text via the reader so the demo still trains."""
import httpx
from app.rag import extract
client = httpx.AsyncClient(transport=_blocked_store_with_reader_transport(),
follow_redirects=True)
text, pages = await extract.crawl_site("https://shop.test/", client=client,
remote_fallback=True)
await client.aclose()
assert "juegos de mesa" in text
assert "Devoluciones" in text
assert pages == ["https://shop.test/"] # the home page is recorded
async def test_no_remote_fallback_by_default_leaves_empty():
"""Without remote_fallback the blocked store yields nothing (no surprise
external calls on the normal path)."""
import httpx
from app.rag import extract
client = httpx.AsyncClient(transport=_blocked_store_with_reader_transport(),
follow_redirects=True)
text, _pages = await extract.crawl_site("https://shop.test/", client=client)
await client.aclose()
assert text.strip() == ""
|