"""image-fetch phase: SSRF guard, fetch+process+store, run completion, resume (§6.5.4/§6.9).""" import io import threading from http.server import BaseHTTPRequestHandler, HTTPServer import psycopg import pytest from PIL import Image from app.domains.products import imagefetch, repo from app.platform import db, objectstore def _png(w, h): buf = io.BytesIO(); Image.new("RGB", (w, h), (1, 2, 3)).save(buf, "PNG"); return buf.getvalue() class _Host(BaseHTTPRequestHandler): GOOD = _png(800, 800) TINY = _png(64, 64) def log_message(self, *a): pass def do_GET(self): if self.path.startswith("/good.png"): self.send_response(200); self.send_header("Content-Type", "image/png") self.end_headers(); self.wfile.write(self.GOOD) elif self.path.startswith("/tiny.png"): self.send_response(200); self.send_header("Content-Type", "image/png") self.end_headers(); self.wfile.write(self.TINY) else: self.send_response(404); self.end_headers() @pytest.fixture() def host(): srv = HTTPServer(("127.0.0.1", 0), _Host) t = threading.Thread(target=srv.serve_forever, daemon=True); t.start() yield f"http://127.0.0.1:{srv.server_address[1]}" srv.shutdown() @pytest.fixture() def pool(fresh_db_url): with psycopg.connect(fresh_db_url) as conn: db.migrate(conn); conn.commit() p = db.open_pool(fresh_db_url, max_size=6) yield p p.close() def _seed(pool): """account+storefront+run; returns (storefront_id, account_id, run_id).""" with pool.connection() as conn: acct = conn.execute("INSERT INTO account (email) VALUES ('m@example.com') RETURNING id").fetchone()[0] sf = conn.execute("INSERT INTO storefront (name) VALUES ('Shop') RETURNING id").fetchone()[0] conn.execute("INSERT INTO storefront_membership (account_id, storefront_id) VALUES (%s,%s)", (acct, sf)) rid = conn.execute( "INSERT INTO import_run (storefront_id, account_id, file_name, dialect," " products_added, products_updated, rows_errored, status)" " VALUES (%s,%s,'c.csv','canonical',0,0,0,'fetching_images') RETURNING id", (sf, acct)).fetchone()[0] conn.commit() return sf, acct, rid def _product(pool, sf, handle): with pool.connection() as conn: pid = conn.execute("INSERT INTO product (storefront_id, handle, title) VALUES (%s,%s,%s) RETURNING id", (sf, handle, handle.title())).fetchone()[0] conn.commit() return pid def _image(pool, pid, url, rid, position=1): with pool.connection() as conn: iid = conn.execute("INSERT INTO product_image (product_id, source_url, position, status, import_run_id)" " VALUES (%s,%s,%s,'pending',%s) RETURNING id", (pid, url, position, rid)).fetchone()[0] conn.commit() return iid def _img_row(pool, iid): with pool.connection() as conn: r = conn.execute("SELECT status, key_original, key_thumb, key_card, key_detail FROM product_image WHERE id=%s", (iid,)).fetchone() return {"status": r[0], "key_original": r[1], "key_thumb": r[2], "key_card": r[3], "key_detail": r[4]} def _run_status(pool, rid): with pool.connection() as conn: return conn.execute("SELECT status FROM import_run WHERE id=%s", (rid,)).fetchone()[0] def test_ssrf_guard_rejects_private_when_disallowed(): with pytest.raises(imagefetch.FetchBlocked): imagefetch.fetch_bytes("http://127.0.0.1:1/x.png", allow_private=False) with pytest.raises(imagefetch.FetchBlocked): imagefetch.fetch_bytes("http://169.254.169.254/latest/meta-data", allow_private=False) with pytest.raises(imagefetch.FetchBlocked): imagefetch.fetch_bytes("ftp://example.com/x", allow_private=False) def test_ssrf_guard_allows_private_when_enabled(host): data, ctype = imagefetch.fetch_bytes(f"{host}/good.png", allow_private=True) assert ctype.startswith("image/") and len(data) > 0 def test_run_phase_classifies_each_image(pool, host, tmp_path): store = objectstore.LocalObjectStore(str(tmp_path)) sf, _acct, rid = _seed(pool) p = _product(pool, sf, "lamp") ok = _image(pool, p, f"{host}/good.png", rid, position=1) _image(pool, p, f"{host}/tiny.png", rid, position=2) _image(pool, p, f"{host}/missing.png", rid, position=3) imagefetch.run_image_phase(pool, store, rid, allow_private=True) with pool.connection() as conn: counts = repo.run_image_counts(conn, rid) assert counts["fetched"] == 1 and counts["rejected"] == 1 and counts["failed"] == 1 row = _img_row(pool, ok) assert row["status"] == "fetched" for k in ("key_original", "key_thumb", "key_card", "key_detail"): assert row[k] and store.get(row[k]) assert _run_status(pool, rid) == "complete_with_problems" def test_unexpected_processing_error_marks_image_failed_not_wedged(pool, host, tmp_path, monkeypatch): # If anything after the fetch raises unexpectedly (e.g. objectstore.put), # the image is marked failed and the run still reaches a terminal status — # never stranded in fetching_images. store = objectstore.LocalObjectStore(str(tmp_path)) sf, _acct, rid = _seed(pool) p = _product(pool, sf, "lamp") iid = _image(pool, p, f"{host}/good.png", rid, position=1) def _boom(*a, **k): raise RuntimeError("storage down") monkeypatch.setattr(store, "put", _boom) imagefetch.run_image_phase(pool, store, rid, allow_private=True) with pool.connection() as conn: counts = repo.run_image_counts(conn, rid) assert counts["failed"] == 1 and counts["pending"] == 0 assert _run_status(pool, rid) == "complete_with_problems" row = _img_row(pool, iid) assert row["status"] == "failed" def test_resume_after_kill_completes_remaining(pool, host, tmp_path): store = objectstore.LocalObjectStore(str(tmp_path)) sf, _acct, rid = _seed(pool) p = _product(pool, sf, "lamp") a = _image(pool, p, f"{host}/good.png", rid, position=1) with pool.connection() as conn: # simulate crash: one already fetched, run still fetching_images repo.mark_image_fetched(conn, a, {"original": "o", "thumb": "t", "card": "c", "detail": "d"}) conn.commit() _image(pool, p, f"{host}/good.png?2", rid, position=2) # still pending resumed = imagefetch.recover_incomplete_runs(pool, store, allow_private=True) assert rid in resumed with pool.connection() as conn: assert repo.run_image_counts(conn, rid)["pending"] == 0 assert _run_status(pool, rid) == "complete"