Validate signature CLI PNG data (#1580)

* Validate signature CLI PNG data

* Keep signature CLI test imports isolated

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 ""
+    png_bytes = _decode_png_data(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}")
     out = Path(args.png)
     out.parent.mkdir(parents=True, exist_ok=True)
     out.write_bytes(png_bytes)

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")