diff --git a/scripts/odysseus-signature b/scripts/odysseus-signature index 1236afa..8cc0e6d 100755 --- a/scripts/odysseus-signature +++ b/scripts/odysseus-signature @@ -29,6 +29,16 @@ except ModuleNotFoundError as e: sys.exit(2) +def _decode_png_data(data_png: str) -> bytes: + raw = data_png or "" + if "," in raw: + raw = raw.split(",", 1)[1] + try: + return base64.b64decode(raw, validate=True) + except Exception as e: + fail(f"data_png is not valid base64: {e}") + + def cmd_list(args): """No `Signature` SQLAlchemy model is registered for the `signatures` table — query via raw SQL so we don't depend on it.""" @@ -85,13 +95,7 @@ def cmd_export(args): ), {"id": args.id}).mappings().first() if not row: fail(f"no signature with id {args.id!r}") - raw = row["data_png"] or "" - if "," in raw: - raw = raw.split(",", 1)[1] - try: - png_bytes = base64.b64decode(raw) - except Exception as e: - fail(f"data_png is not valid base64: {e}") + png_bytes = _decode_png_data(row["data_png"] or "") out = Path(args.png) out.parent.mkdir(parents=True, exist_ok=True) out.write_bytes(png_bytes) diff --git a/tests/test_signature_cli_export.py b/tests/test_signature_cli_export.py new file mode 100644 index 0000000..ffc9757 --- /dev/null +++ b/tests/test_signature_cli_export.py @@ -0,0 +1,40 @@ +import importlib.machinery +import importlib.util +import sys +from pathlib import Path +from types import ModuleType + + +def _load_signature_cli(monkeypatch): + sqlalchemy_mod = ModuleType("sqlalchemy") + sqlalchemy_mod.text = lambda value: value + core_mod = ModuleType("core") + database_mod = ModuleType("core.database") + database_mod.engine = object() + monkeypatch.setitem(sys.modules, "sqlalchemy", sqlalchemy_mod) + monkeypatch.setitem(sys.modules, "core", core_mod) + monkeypatch.setitem(sys.modules, "core.database", database_mod) + + path = Path(__file__).resolve().parent.parent / "scripts" / "odysseus-signature" + loader = importlib.machinery.SourceFileLoader("odysseus_signature_cli_under_test", str(path)) + spec = importlib.util.spec_from_loader(loader.name, loader) + module = importlib.util.module_from_spec(spec) + loader.exec_module(module) + return module + + +def test_decode_png_data_accepts_data_url(monkeypatch): + cli = _load_signature_cli(monkeypatch) + + assert cli._decode_png_data("data:image/png;base64,aGVsbG8=") == b"hello" + + +def test_decode_png_data_rejects_invalid_base64(monkeypatch): + cli = _load_signature_cli(monkeypatch) + + try: + cli._decode_png_data("not valid!!!") + except SystemExit as exc: + assert exc.code == 1 + else: + raise AssertionError("expected invalid base64 to exit")