206 lines
7.2 KiB
Python
206 lines
7.2 KiB
Python
import json
|
|
import tempfile
|
|
import zipfile
|
|
import zlib
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
from aiohttp import web
|
|
from aiohttp.test_utils import make_mocked_request
|
|
from mjr_am_backend.routes.handlers import batch_zip as bz
|
|
from mjr_am_backend.shared import Result
|
|
|
|
|
|
def _app():
|
|
app = web.Application()
|
|
routes = web.RouteTableDef()
|
|
bz.register_batch_zip_routes(routes)
|
|
app.add_routes(routes)
|
|
return app
|
|
|
|
|
|
def _png_chunk(kind: bytes, data: bytes = b"") -> bytes:
|
|
return (
|
|
len(data).to_bytes(4, "big")
|
|
+ kind
|
|
+ data
|
|
+ zlib.crc32(kind + data).to_bytes(4, "big")
|
|
)
|
|
|
|
|
|
def _png_with_workflow_text() -> bytes:
|
|
return b"\x89PNG\r\n\x1a\n" + _png_chunk(b"tEXt", b"workflow\x00secret") + _png_chunk(b"IEND")
|
|
|
|
|
|
def test_batch_zip_sanitize_and_cleanup_helpers(monkeypatch):
|
|
with tempfile.TemporaryDirectory() as td:
|
|
tmp_path = Path(td)
|
|
monkeypatch.setattr(bz, "_BATCH_CACHE", {})
|
|
monkeypatch.setattr(bz, "_BATCH_MAX", 1)
|
|
monkeypatch.setattr(bz, "_BATCH_TTL_SECONDS", 0)
|
|
monkeypatch.setattr(bz, "_BATCH_DIR", tmp_path)
|
|
assert bz._sanitize_token("short") == ""
|
|
token = "a" * 40
|
|
assert bz._sanitize_token(token) == token
|
|
|
|
z = tmp_path / "x.zip"
|
|
z.write_bytes(b"x")
|
|
bz._BATCH_CACHE[token] = {"path": z, "created_at": 1.0}
|
|
bz._cleanup_batch_zips()
|
|
assert token not in bz._BATCH_CACHE
|
|
|
|
|
|
def test_batch_file_response_sanitizes_content_disposition(tmp_path: Path) -> None:
|
|
token = "a" * 40
|
|
zip_path = tmp_path / "payload.zip"
|
|
zip_path.write_bytes(b"zip")
|
|
entry = {"path": zip_path, "filename": 'bad\r\nname";\x01.zip'}
|
|
response = bz._batch_file_response(entry, token)
|
|
header = str(response.headers.get("Content-Disposition") or "")
|
|
assert "\r" not in header
|
|
assert "\n" not in header
|
|
assert '"' in header
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_batch_zip_create_invalid_and_get_not_found(monkeypatch):
|
|
app = _app()
|
|
monkeypatch.setattr(bz, "_csrf_error", lambda _r: None)
|
|
monkeypatch.setattr(bz, "_check_rate_limit", lambda *args, **kwargs: (True, None))
|
|
monkeypatch.setattr(bz, "_require_write_access", lambda _r: Result.Ok({}))
|
|
|
|
async def _read_json(_request, max_bytes=None):
|
|
_ = max_bytes
|
|
return Result.Ok({"token": "bad", "items": []})
|
|
|
|
monkeypatch.setattr(bz, "_read_json", _read_json)
|
|
|
|
req = make_mocked_request("POST", "/mjr/am/batch-zip", app=app)
|
|
match = await app.router.resolve(req)
|
|
resp = await match.handler(req)
|
|
body = json.loads(resp.text)
|
|
assert body.get("code") == "INVALID_INPUT"
|
|
|
|
req2 = make_mocked_request("GET", "/mjr/am/batch-zip/" + ("a" * 40), app=app)
|
|
req2._match_info = {"token": "a" * 40}
|
|
match2 = await app.router.resolve(req2)
|
|
resp2 = await match2.handler(req2)
|
|
body2 = json.loads(resp2.text)
|
|
assert body2.get("code") in {"NOT_FOUND", "NOT_READY"}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_batch_zip_create_and_fetch_ready(monkeypatch, tmp_path: Path):
|
|
app = _app()
|
|
monkeypatch.setattr(bz, "_csrf_error", lambda _r: None)
|
|
monkeypatch.setattr(bz, "_check_rate_limit", lambda *args, **kwargs: (True, None))
|
|
monkeypatch.setattr(bz, "_require_write_access", lambda _r: Result.Ok({}))
|
|
monkeypatch.setattr(bz, "_BATCH_DIR", tmp_path)
|
|
monkeypatch.setattr(bz, "_BATCH_CACHE", {})
|
|
|
|
src = tmp_path / "a.png"
|
|
src.write_bytes(b"data")
|
|
token = "b" * 40
|
|
|
|
async def _read_json(_request, max_bytes=None):
|
|
_ = max_bytes
|
|
return Result.Ok({"token": token, "items": [{"filename": "a.png", "subfolder": "", "type": "output"}]})
|
|
|
|
monkeypatch.setattr(bz, "_read_json", _read_json)
|
|
monkeypatch.setattr(bz, "_resolve_item_path", lambda _item: (src, tmp_path))
|
|
|
|
req = make_mocked_request("POST", "/mjr/am/batch-zip", app=app)
|
|
match = await app.router.resolve(req)
|
|
resp = await match.handler(req)
|
|
body = json.loads(resp.text)
|
|
assert body.get("ok") is True
|
|
|
|
req2 = make_mocked_request("GET", f"/mjr/am/batch-zip/{token}", app=app)
|
|
req2._match_info = {"token": token}
|
|
match2 = await app.router.resolve(req2)
|
|
resp2 = await match2.handler(req2)
|
|
assert getattr(resp2, "status", 200) in {200, 206}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_batch_zip_create_clean_strips_png_metadata(monkeypatch, tmp_path: Path):
|
|
app = _app()
|
|
monkeypatch.setattr(bz, "_csrf_error", lambda _r: None)
|
|
monkeypatch.setattr(bz, "_check_rate_limit", lambda *args, **kwargs: (True, None))
|
|
monkeypatch.setattr(bz, "_require_write_access", lambda _r: Result.Ok({}))
|
|
monkeypatch.setattr(bz, "_BATCH_DIR", tmp_path)
|
|
monkeypatch.setattr(bz, "_BATCH_CACHE", {})
|
|
|
|
src = tmp_path / "a.png"
|
|
src.write_bytes(_png_with_workflow_text())
|
|
token = "d" * 40
|
|
|
|
async def _read_json(_request, max_bytes=None):
|
|
_ = max_bytes
|
|
return Result.Ok(
|
|
{
|
|
"token": token,
|
|
"strip_metadata": True,
|
|
"items": [{"filename": "a.png", "subfolder": "", "type": "output"}],
|
|
}
|
|
)
|
|
|
|
monkeypatch.setattr(bz, "_read_json", _read_json)
|
|
monkeypatch.setattr(bz, "_resolve_item_path", lambda _item: (src, tmp_path))
|
|
|
|
req = make_mocked_request("POST", "/mjr/am/batch-zip", app=app)
|
|
match = await app.router.resolve(req)
|
|
resp = await match.handler(req)
|
|
body = json.loads(resp.text)
|
|
assert body.get("ok") is True
|
|
assert body["data"]["filename"] == "Majoor_Clean_Batch_1.zip"
|
|
|
|
entry = bz._BATCH_CACHE[token]
|
|
with zipfile.ZipFile(entry["path"], "r") as zf:
|
|
cleaned = zf.read("a.png")
|
|
|
|
assert b"workflow\x00secret" not in cleaned
|
|
assert cleaned.startswith(b"\x89PNG\r\n\x1a\n")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_batch_zip_rejects_inode_swap_to_mitigate_toctou(monkeypatch, tmp_path: Path):
|
|
app = _app()
|
|
monkeypatch.setattr(bz, "_csrf_error", lambda _r: None)
|
|
monkeypatch.setattr(bz, "_check_rate_limit", lambda *args, **kwargs: (True, None))
|
|
monkeypatch.setattr(bz, "_require_write_access", lambda _r: Result.Ok({}))
|
|
monkeypatch.setattr(bz, "_BATCH_DIR", tmp_path)
|
|
monkeypatch.setattr(bz, "_BATCH_CACHE", {})
|
|
|
|
src = tmp_path / "toctou.png"
|
|
src.write_bytes(b"data")
|
|
token = "c" * 40
|
|
|
|
async def _read_json(_request, max_bytes=None):
|
|
_ = max_bytes
|
|
return Result.Ok({"token": token, "items": [{"filename": "toctou.png", "subfolder": "", "type": "output"}]})
|
|
|
|
monkeypatch.setattr(bz, "_read_json", _read_json)
|
|
monkeypatch.setattr(bz, "_resolve_item_path", lambda _item: (src, tmp_path))
|
|
|
|
real_fstat = bz.os.fstat
|
|
|
|
def _fstat_inode_mismatch(fd):
|
|
st = real_fstat(fd)
|
|
return SimpleNamespace(
|
|
st_mode=int(getattr(st, "st_mode", 0)),
|
|
st_mtime=float(getattr(st, "st_mtime", 0.0)),
|
|
st_size=int(getattr(st, "st_size", 0)),
|
|
st_ino=int(getattr(st, "st_ino", 0)) + 1,
|
|
)
|
|
|
|
monkeypatch.setattr(bz.os, "fstat", _fstat_inode_mismatch)
|
|
|
|
req = make_mocked_request("POST", "/mjr/am/batch-zip", app=app)
|
|
match = await app.router.resolve(req)
|
|
resp = await match.handler(req)
|
|
body = json.loads(resp.text)
|
|
assert body.get("ok") is False
|
|
assert body.get("code") == "NO_VALID_FILES"
|