fix: drop null vault keys, cap snapshot restore size, isolate test globals

Vaults.add_keys and kv add skip null content keys instead of failing the batch; is_null_key moves to core.vault. A manifest snapshot restore has one 32 MiB inflate budget. Tests reset the speed limiter and config between runs.

docs/reference/configuration/vaults.md

@@ -36,8 +36,8 @@ key_vaults:
     - A vault of `type: API` whose `name` contains `decrypt_labs` auto-fills its `token` from
       [`decrypt_labs_api_key`](misc.md#external-api-keys) when not set inline. Vault `type` values are
       case-sensitive module names.
-    - unshackle treats an all-zero content key (32 zeros) as "no key" everywhere and never
-      stores it.
+    - unshackle treats an empty or all-zero content key, of any length, as "no key" everywhere
+      and never stores it.
 
 !!! note "Bad content keys"
     A content key from a vault is proven by short FFmpeg decodes of the decrypted output,

tests/commands/test_kv_vault_path.py

@@ -36,3 +36,19 @@ def test_kv_copy_rejects_a_name_that_is_neither_configured_nor_a_file(monkeypatc
 
     assert result.exit_code != 0
     assert "not defined in the config" in result.output
+
+
+def test_kv_add_skips_a_null_key_line_and_adds_the_rest(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
+    keys = tmp_path / "keys.txt"
+    null_kid = "0" * 31 + "2"
+    keys.write_text(f"{KID}:{KEY}\n{null_kid}:{'0' * 32}\n", encoding="utf8")
+    dest = tmp_path / "local.db"
+    monkeypatch.setattr(kv_module.config, "key_vaults", [{"type": "SQLite", "name": "local", "path": str(dest)}])
+    monkeypatch.setattr(kv_module.Services, "get_vault_tag", lambda name: name)
+
+    result = CliRunner().invoke(kv_module.kv, ["add", str(keys), "TESTSVC", "local"])
+
+    assert result.exit_code == 0, result.output
+    check = SQLite("check", dest)
+    assert check.get_key(KID, "TESTSVC") == KEY
+    assert check.get_key(null_kid, "TESTSVC") is None

tests/conftest.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import importlib
 import socket
 import socketserver
 from pathlib import Path
@@ -14,6 +15,8 @@ from unshackle.core.config import config
 from unshackle.core.proxies.proxy import CATALOGUES
 from unshackle.core.update_checker import UpdateChecker
 
+requests_downloader = importlib.import_module("unshackle.core.downloaders.requests")
+
 # shutdown() blocks until the serve loop polls again, and the 0.5s default is paid by every
 # test server at teardown. The default changes for the whole process, servers in core included.
 socketserver.BaseServer.serve_forever.__defaults__ = (0.01,)
@@ -59,6 +62,13 @@ def default_codec_names(monkeypatch: pytest.MonkeyPatch) -> None:
     monkeypatch.setattr(config, "codec_names", {})
 
 
+@pytest.fixture(autouse=True)
+def no_speed_limit(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Start with no speed limit, so a limit that an earlier test locked cannot block a later one."""
+    monkeypatch.setattr(requests_downloader, "speed_limiter", None)
+    monkeypatch.setattr(requests_downloader, "speed_limit_locked", False)
+
+
 @pytest.fixture
 def mux_args(monkeypatch: pytest.MonkeyPatch) -> Callable[..., list[str]]:
     """Run Tracks.mux up to mkvmerge and return its arguments; ``ids`` is what the track ID probe reports."""

tests/core/test_manifest_snapshot.py

@@ -2,13 +2,17 @@
 
 from __future__ import annotations
 
+import base64
 import json
+import logging
+import zlib
 from copy import deepcopy
 from typing import Any, Optional
 
 import pytest
 
 from unshackle.core.manifests import DASH
+from unshackle.core.manifests import snapshot as snapshot_module
 from unshackle.core.manifests.snapshot import capture, parse, restore
 from unshackle.core.tracks import Tracks
 from unshackle.core.tracks.track import Track
@@ -117,6 +121,10 @@ def round_trip(tracks: Tracks, level: int = -1) -> Tracks:
     return client
 
 
+def inflated_size(entry: dict[str, Any]) -> int:
+    return len(zlib.decompress(base64.b64decode(entry["data"])))
+
+
 def test_service_filter_drops_every_ad_period_after_the_first_content_period():
     assert stitched(service_tracks().videos[0]) == ["main-1", "main-2", "main-3"]
 
@@ -176,6 +184,34 @@ def test_a_saved_manifest_that_does_not_parse_gives_none():
     assert parse({"manifests": [{"type": "dash", "url": URL, "data": "bm90IHpsaWI="}]}, URL, "en") is None
 
 
+def test_the_inflate_budget_covers_every_entry_of_one_restore(monkeypatch, caplog):
+    saved = capture(service_tracks())
+    entry = saved["manifests"][0]
+    size = inflated_size(entry)
+    saved["manifests"].append({**entry, "url": URL + "?copy"})
+    monkeypatch.setattr(snapshot_module, "MAX_SNAPSHOT_BYTES", size)
+    assert parse(saved, URL, "en") is not None
+
+    client = Tracks()
+    client.add(Track.from_dict(service_tracks().videos[0].to_dict()))
+    with caplog.at_level(logging.WARNING, logger="snapshot"):
+        restore(client, saved)
+    assert client.videos[0].data.get("dash")
+    assert ["expands beyond allowed size" in r.getMessage() for r in caplog.records] == [True]
+
+    monkeypatch.setattr(snapshot_module, "MAX_SNAPSHOT_BYTES", size - 1)
+    assert parse(saved, URL, "en") is None
+
+
+def test_the_fragments_and_the_whole_manifests_share_one_budget(monkeypatch):
+    tracks = service_tracks()
+    detach_video(tracks)
+    monkeypatch.setattr(snapshot_module, "MAX_SNAPSHOT_BYTES", inflated_size(capture(tracks)["manifests"][0]))
+    client = round_trip(tracks)
+    assert client.videos[0].data.get("dash")
+    assert not client.audio[0].data.get("dash")
+
+
 def test_compression_level_zero_restores():
     video = round_trip(service_tracks(), level=0).videos[0]
     assert stitched(video) == ["main-1", "main-2", "main-3"]

tests/core/test_vault_table_case.py

@@ -17,7 +17,8 @@ from uuid import UUID
 
 import pytest
 
-from unshackle.core.vault import Vault
+from unshackle.core.vault import Vault, is_null_key
+from unshackle.core.vaults import Vaults
 from unshackle.vaults.SQLite import SQLite
 
 LOG = logging.getLogger("kv-test")
@@ -119,10 +120,29 @@ def test_synced_lowercase_table_survives_the_service_being_added_later(tmp_path:
     assert len(list(v.get_keys("SVCC"))) == 2
 
 
+@pytest.mark.parametrize(
+    ("key", "expected"),
+    [
+        (None, True),
+        ("", True),
+        ("0" * 32, True),
+        ("00", True),
+        ("00112233445566778899aabbccddeeff", False),
+        ("00112233445566778899AABBCCDDEEFF", False),
+    ],
+)
+def test_is_null_key(key: Optional[str], expected: bool) -> None:
+    assert is_null_key(key) is expected
+
+
 def test_null_keys_are_rejected_and_filtered(tmp_path: Path) -> None:
     v = SQLite("test", tmp_path / "kv.db")
-    with pytest.raises(ValueError):
-        v.add_keys("SVCA", {KID: "0" * 32})
+    null_key_error = r"^You cannot add a NULL Content Key to a Vault\.$"
+    for key in ("0" * 32, ""):
+        with pytest.raises(ValueError, match=null_key_error):
+            v.add_keys("SVCA", {KID: key})
+        with pytest.raises(ValueError, match=null_key_error):
+            v.add_key("SVCA", KID, key)
 
     v.create_table("svca")
     conn = sqlite3.connect(v.path)
@@ -134,6 +154,15 @@ def test_null_keys_are_rejected_and_filtered(tmp_path: Path) -> None:
     assert list(v.get_keys("SVCA")) == []
 
 
+def test_vaults_push_drops_null_keys_and_keeps_the_rest(tmp_path: Path) -> None:
+    vaults = Vaults("SVCA")
+    vaults.vaults = [SQLite("test", tmp_path / "kv.db")]
+    other_kid = "0" * 31 + "2"
+    assert vaults.add_keys({KID: KEY, other_kid: "0" * 32}) == 1
+    assert vaults.vaults[0].get_key(KID, "SVCA") == KEY
+    assert vaults.vaults[0].get_key(other_kid, "SVCA") is None
+
+
 def test_get_services_yields_stored_table_names(vault: SQLite, aliased_service: str) -> None:
     """Names are yielded verbatim. Normalising the `svcdalias` table to the installed SVCD
     tag would hand back a name that addresses no table at all."""

tests/downloaders/test_multiprocess.py

@@ -180,11 +180,8 @@ def test_speed_limit_forces_single_process(server, tmp_path, monkeypatch):
     n = dl.MP_MIN_SEGMENTS
     urls = [{"url": url(server, f"/seg/{i}.bin")} for i in range(n)]
     dl.set_speed_limit(1_000_000_000)  # far above the tiny batch, so the test is not slowed
-    try:
-        for _ in dl.requests(urls, output_dir=tmp_path, filename="seg_{i:04}.bin", max_workers=4, processes=2):
-            pass
-    finally:
-        dl.set_speed_limit(None)
+    for _ in dl.requests(urls, output_dir=tmp_path, filename="seg_{i:04}.bin", max_workers=4, processes=2):
+        pass
 
     files = sorted(tmp_path.glob("seg_*.bin"))
     assert len(files) == n

tests/downloaders/test_speed_limit.py

@@ -121,15 +121,11 @@ def test_bucket_chunk_larger_than_capacity_does_not_deadlock(monkeypatch):
 
 def test_locked_limit_ignores_unlocked_calls():
     """serve's global_speed_limit must survive the per-job set in dl.result()."""
-    try:
-        set_speed_limit(5_000_000, lock=True)
-        set_speed_limit(None)
-        set_speed_limit(1_000_000)
-        limiter = requests_downloader.speed_limiter
-        assert limiter is not None and limiter.rate == 5_000_000
-    finally:
-        requests_downloader.speed_limit_locked = False
-        requests_downloader.speed_limiter = None
+    set_speed_limit(5_000_000, lock=True)
+    set_speed_limit(None)
+    set_speed_limit(1_000_000)
+    limiter = requests_downloader.speed_limiter
+    assert limiter is not None and limiter.rate == 5_000_000
 
 
 def test_bucket_wait_aborts_on_cancellation():

tests/remote/unit/test_serve_reload.py

@@ -26,6 +26,9 @@ def running(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> tuple[web.Applic
     wvds.mkdir()
     prds.mkdir()
     (wvds / "dev.wvd").write_bytes(b"x")
+    # a valid reload replaces every config section, so put each one back for the next test
+    for name, value in list(vars(config).items()):
+        monkeypatch.setattr(config, name, value)
     monkeypatch.setattr(config.directories, "wvds", wvds)
     monkeypatch.setattr(config.directories, "prds", prds)
     monkeypatch.setattr(config, "serve", {"api_secret": "api", "users": {"old": {"devices": ["dev"]}}})

tests/remote/unit/test_worker_speed_limit.py

@@ -15,16 +15,6 @@ requests_downloader = importlib.import_module("unshackle.core.downloaders.reques
 pytestmark = pytest.mark.unit
 
 
-@pytest.fixture(autouse=True)
-def clean_limiter():
-    """Start and end each test with no limiter, whatever an earlier test leaked."""
-    requests_downloader.speed_limit_locked = False
-    requests_downloader.speed_limiter = None
-    yield
-    requests_downloader.speed_limit_locked = False
-    requests_downloader.speed_limiter = None
-
-
 @pytest.mark.parametrize(("lock", "expected"), [(True, 2_000_000.0), (False, None)])
 def test_worker_share_of_limit(lock, expected):
     requests_downloader.set_speed_limit(8_000_000, lock=lock)

unshackle/commands/kv.py

@@ -2,7 +2,7 @@ import logging
 import re
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from pathlib import Path
-from typing import Optional, Union, cast
+from typing import Iterable, Optional, Union, cast
 from uuid import UUID
 
 import click
@@ -23,7 +23,7 @@ from unshackle.core.config import config
 from unshackle.core.console import GradientPulseBarColumn, console
 from unshackle.core.constants import context_settings
 from unshackle.core.services import Services
-from unshackle.core.vault import Vault
+from unshackle.core.vault import Vault, is_null_key
 from unshackle.core.vaults import Vaults
 
 
@@ -48,15 +48,20 @@ def load_vaults(vault_names: list[str]) -> Vaults:
     return vaults
 
 
+def without_null_keys(pairs: Iterable[tuple[str, str]], log: logging.Logger) -> dict[str, str]:
+    """Return the (kid, key) pairs that are not null as a dict, with a warning for each one dropped."""
+    kid_keys: dict[str, str] = {}
+    for kid, key in pairs:
+        if is_null_key(key):
+            log.warning(f"Skipping NULL key: {kid}:{key}")
+        else:
+            kid_keys[kid] = key
+    return kid_keys
+
+
 def process_service_keys(from_vault: Vault, service: str, log: logging.Logger) -> dict[str, str]:
     """Get and validate keys from a vault for a specific service."""
-    content_keys = list(from_vault.get_keys(service))
-
-    bad_keys = {kid: key for kid, key in content_keys if not key or key.count("0") == len(key)}
-    for kid, key in bad_keys.items():
-        log.warning(f"Skipping NULL key: {kid}:{key}")
-
-    return {kid: key for kid, key in content_keys if kid not in bad_keys}
+    return without_null_keys(from_vault.get_keys(service), log)
 
 
 class _PaddedProgress(Progress):
@@ -301,15 +306,10 @@ def add(file: Path, service: str, vaults: list[str]) -> None:
     vaults_ = load_vaults(list(vaults))
 
     data = file.read_text(encoding="utf8")
-    kid_keys: dict[str, str] = {}
-    for line in data.splitlines(keepends=False):
-        line = line.strip()
-        match = re.search(r"^(?P<kid>[0-9a-fA-F]{32}):(?P<key>[0-9a-fA-F]{32})$", line)
-        if not match:
-            continue
-        kid = match.group("kid").lower()
-        key = match.group("key").lower()
-        kid_keys[kid] = key
+    matches = (
+        re.search(r"^(?P<kid>[0-9a-fA-F]{32}):(?P<key>[0-9a-fA-F]{32})$", line.strip()) for line in data.splitlines()
+    )
+    kid_keys = without_null_keys(((m["kid"].lower(), m["key"].lower()) for m in matches if m), log)
 
     total_count = len(kid_keys)
 
@@ -331,7 +331,7 @@ def search_vault(vault: Vault, kid: str, services: list[str], log: logging.Logge
         except Exception as e:
             log.debug(f"{vault} [{svc}]: lookup error ({e})")
             return None
-        if key and key.count("0") != len(key):
+        if key and not is_null_key(key):
             return svc, key
         return None
 

unshackle/core/cdm/custom_remote_cdm.py

@@ -11,6 +11,7 @@ from pywidevine.device import DeviceTypes
 from requests import Session
 
 from unshackle.core import __version__
+from unshackle.core.vault import is_null_key
 from unshackle.core.vaults import Vaults
 
 
@@ -812,7 +813,7 @@ class CustomRemoteCDM:
                     else:
                         kid_uuid = UUID(hex=clean_kid.ljust(32, "0"))
                     key, _ = self.vaults.get_key(kid_uuid)
-                    if key and key.count("0") != len(key):
+                    if not is_null_key(key):
                         vault_keys.append({"kid": kid_str, "key": key, "type": "CONTENT"})
                 except (ValueError, TypeError):
                     continue

unshackle/core/cdm/decrypt_labs_remote_cdm.py

@@ -11,6 +11,7 @@ from pywidevine.device import DeviceTypes
 from requests import Session
 
 from unshackle.core import __version__
+from unshackle.core.vault import is_null_key
 from unshackle.core.vaults import Vaults
 
 
@@ -412,7 +413,7 @@ class DecryptLabsRemoteCDM:
                     else:
                         kid_uuid = UUID(hex=clean_kid.ljust(32, "0"))
                     key, _ = self.vaults.get_key(kid_uuid)
-                    if key and key.count("0") != len(key):
+                    if not is_null_key(key):
                         vault_keys.append({"kid": kid_str, "key": key, "type": "CONTENT"})
                 except (ValueError, TypeError):
                     continue

unshackle/core/drm/__init__.py

@@ -1,5 +1,5 @@
 import base64
-from typing import Any, Optional, Union
+from typing import Any, Union
 from uuid import UUID
 
 from unshackle.core.drm.blowfish_stripe import BlowfishStripe
@@ -8,6 +8,7 @@ from unshackle.core.drm.clearkey_cenc import ClearKeyCENC
 from unshackle.core.drm.monalisa import MonaLisa
 from unshackle.core.drm.playready import PlayReady
 from unshackle.core.drm.widevine import Widevine
+from unshackle.core.vault import is_null_key
 
 # BlowfishStripe stays out: the licensing paths typed on DRM_T never receive a DRM that needs no licence
 DRM_T = Union[ClearKey, ClearKeyCENC, Widevine, PlayReady, MonaLisa]
@@ -50,11 +51,6 @@ def drm_from_dict(data: dict[str, Any]) -> Union[Widevine, PlayReady, ClearKeyCE
     return drm
 
 
-def is_null_key(key: Optional[str]) -> bool:
-    """True for a blank or all-zero content key, the null key that no vault stores."""
-    return not key or key.count("0") == len(key)
-
-
 def drop_null_keys(drm: Any) -> None:
     """Remove the null keys from a DRM object's content keys. A read-only key view stays as it is."""
     content_keys = getattr(drm, "content_keys", None)

unshackle/core/manifests/snapshot.py

@@ -24,6 +24,9 @@ log = logging.getLogger("snapshot")
 
 Snapshot = dict[str, list[dict[str, Any]]]
 
+# A peer sets the entry count, so one restore or parse is capped in total. The parsed tree takes about 30x the XML.
+MAX_SNAPSHOT_BYTES = 32 * 1024 * 1024
+
 
 def capture(tracks: Tracks, level: int = zlib.Z_DEFAULT_COMPRESSION) -> Snapshot:
     """The snapshot of these tracks, JSON-safe, at zlib ``level`` (0 gives an uncompressed zlib stream)."""
@@ -36,8 +39,9 @@ def restore(tracks: Tracks, snapshot: Any, language: Any = None) -> None:
 
     ``language`` is the parse fallback when no track has a language of its own.
     """
-    restore_fragments(tracks, entries(snapshot, "track_manifests"))
-    restore_manifests(tracks, entries(snapshot, "manifests"), language)
+    budget = [MAX_SNAPSHOT_BYTES]
+    restore_fragments(tracks, entries(snapshot, "track_manifests"), budget)
+    restore_manifests(tracks, entries(snapshot, "manifests"), budget, language)
 
 
 def parse(snapshot: Any, url: str, language: Any = None) -> Optional[Tracks]:
@@ -53,7 +57,7 @@ def parse(snapshot: Any, url: str, language: Any = None) -> Optional[Tracks]:
     if entry is None:
         return None
     try:
-        tracks = to_tracks(entry, language)
+        tracks = to_tracks(entry, language, [MAX_SNAPSHOT_BYTES])
     except MissingLanguageError:
         raise
     except Exception as e:
@@ -78,10 +82,13 @@ def pack(element: Any, level: int) -> str:
     return base64.b64encode(zlib.compress(xml_bytes, level)).decode("ascii")
 
 
-def unpack(data: str) -> Any:
+def unpack(data: str, budget: list[int]) -> Any:
+    """The saved tree. ``budget[0]`` is the inflated byte count still allowed; this call takes its share."""
     from unshackle.core.api.compression import safe_inflate
 
-    return etree.fromstring(safe_inflate(base64.b64decode(data)))
+    xml_bytes = safe_inflate(base64.b64decode(data), budget[0])
+    budget[0] -= len(xml_bytes)
+    return etree.fromstring(xml_bytes)
 
 
 def counted(manifest: Any, tag: str, counts: dict[int, int]) -> int:
@@ -184,7 +191,7 @@ def fragments(tracks: Tracks, level: int, counts: dict[int, int]) -> list[dict[s
     return out
 
 
-def restore_fragments(tracks: Tracks, saved: list[dict[str, Any]]) -> None:
+def restore_fragments(tracks: Tracks, saved: list[dict[str, Any]], budget: list[int]) -> None:
     """Fill ``track.data`` from the one-track fragments, by track id. This runs before the whole-manifest restore."""
     by_id = {str(t.id): t for t in tracks}
     for entry in saved:
@@ -192,7 +199,7 @@ def restore_fragments(tracks: Tracks, saved: list[dict[str, Any]]) -> None:
         if track is None or entry.get("type") != "dash" or not isinstance(entry.get("data"), str):
             continue
         try:
-            manifest = unpack(entry["data"])
+            manifest = unpack(entry["data"], budget)
             period = manifest.find("Period")
             adaptation_set = period.find("AdaptationSet")
             representation = adaptation_set.find("Representation")
@@ -212,9 +219,9 @@ def restore_fragments(tracks: Tracks, saved: list[dict[str, Any]]) -> None:
         track.descriptor = Track.Descriptor.DASH
 
 
-def to_tracks(entry: dict[str, Any], language: Any) -> Optional[Tracks]:
+def to_tracks(entry: dict[str, Any], language: Any, budget: list[int]) -> Optional[Tracks]:
     """The tracks of one whole-manifest entry, with the service's filtered periods dropped again."""
-    root = unpack(entry["data"])
+    root = unpack(entry["data"], budget)
     if entry.get("type") == "dash":
         skip = set(entry.get("filtered_period_ids") or [])
         return DASH(root, entry["url"]).to_tracks(
@@ -229,7 +236,7 @@ def has_data(track: Any) -> bool:
     return any(track.data.get(key) for key in ("dash", "ism"))
 
 
-def restore_manifests(tracks: Tracks, saved: list[dict[str, Any]], fallback: Any = None) -> None:
+def restore_manifests(tracks: Tracks, saved: list[dict[str, Any]], budget: list[int], fallback: Any = None) -> None:
     """Re-parse each whole manifest and copy ``data`` to the tracks that match, by id or by attributes."""
     targets = [t for t in tracks if not has_data(t)]
     if not targets:
@@ -240,7 +247,7 @@ def restore_manifests(tracks: Tracks, saved: list[dict[str, Any]], fallback: Any
         if not isinstance(entry.get("data"), str) or not isinstance(entry.get("url"), str):
             continue
         try:
-            local = to_tracks(entry, language)
+            local = to_tracks(entry, language, budget)
         except Exception as e:
             log.warning("Failed to re-parse %s manifest from %s: %s", entry.get("type"), entry.get("url"), e)
             continue

unshackle/core/vault.py

@@ -23,6 +23,11 @@ class RateLimited(Exception):
 UNREACHABLE = (RequestsConnectionError, Timeout, ConnectionError, TimeoutError)
 
 
+def is_null_key(key: Optional[str]) -> bool:
+    """True for a blank or all-zero content key, the null key that no vault stores."""
+    return not key or key.count("0") == len(key)
+
+
 def vault_session(timeout: float) -> Session:
     """An HTTP session for a remote vault that retries a 429 and nothing else.
 

unshackle/core/vaults.py

@@ -8,7 +8,7 @@ from uuid import UUID
 
 from unshackle.core.config import config
 from unshackle.core.utilities import get_debug_logger, import_module_by_path
-from unshackle.core.vault import UNREACHABLE, RateLimited, Vault
+from unshackle.core.vault import UNREACHABLE, RateLimited, Vault, is_null_key
 
 log = logging.getLogger(__name__)
 
@@ -136,7 +136,7 @@ class Vaults:
                     duration_ms=round((time.monotonic() - start) * 1000, 1),
                 )
             return None
-        found = bool(key and key.count("0") != len(key))
+        found = not is_null_key(key)
         if found and key and self.is_flagged(kid, key):
             found = False
             if (vault.name, kid, key) not in self.warned:
@@ -236,10 +236,12 @@ class Vaults:
         This method also skips Vaults with no_push=True.
 
         This method pushes to all Vaults at the same time, so one unreachable Vault costs only its
-        own timeout and does not delay the others.
+        own timeout and does not delay the others. This method drops null keys, because a Vault
+        rejects the whole batch when it contains one.
         """
+        kid_keys = {kid: key for kid, key in kid_keys.items() if not is_null_key(key)}
         vaults = [vault for vault in self.live() if not vault.no_push]
-        if not vaults:
+        if not vaults or not kid_keys:
             return 0
         with ThreadPoolExecutor(len(vaults), thread_name_prefix="vault-push") as pool:
             return sum(pool.map(lambda vault: self.push_keys(vault, kid_keys), vaults))

unshackle/vaults/HTTP.py

@@ -3,7 +3,7 @@ from enum import Enum
 from typing import Iterator, Optional, Union
 from uuid import UUID
 
-from unshackle.core.vault import RateLimited, Vault, vault_session
+from unshackle.core.vault import RateLimited, Vault, is_null_key, vault_session
 
 
 class InsertResult(Enum):
@@ -218,7 +218,7 @@ class HTTP(Vault):
                 yield key_entry["kid"], key_entry["key"]
 
     def add_key(self, service: str, kid: Union[UUID, str], key: str) -> bool:
-        if not key or key.count("0") == len(key):
+        if is_null_key(key):
             raise ValueError("You cannot add a NULL Content Key to a Vault.")
 
         if self.api_mode == "decrypt_labs":
@@ -268,7 +268,7 @@ class HTTP(Vault):
             return 0
 
         for kid, key in kid_keys.items():
-            if not key or key.count("0") == len(key):
+            if is_null_key(key):
                 raise ValueError("You cannot add a NULL Content Key to a Vault.")
 
         processed_kid_keys = {
@@ -374,7 +374,7 @@ class HTTP(Vault):
         This method provides more granular feedback than the standard add_key method.
         Available in both API modes.
         """
-        if not key or key.count("0") == len(key):
+        if is_null_key(key):
             raise ValueError("You cannot add a NULL Content Key to a Vault.")
 
         if self.api_mode == "decrypt_labs":

unshackle/vaults/MySQL.py

@@ -8,7 +8,7 @@ import pymysql
 from pymysql.cursors import DictCursor
 
 from unshackle.core.config import config
-from unshackle.core.vault import Vault
+from unshackle.core.vault import Vault, is_null_key
 
 
 class MySQL(Vault):
@@ -74,7 +74,7 @@ class MySQL(Vault):
             cursor.close()
 
     def add_key(self, service: str, kid: Union[UUID, str], key: str) -> bool:
-        if not key or key.count("0") == len(key):
+        if is_null_key(key):
             raise ValueError("You cannot add a NULL Content Key to a Vault.")
 
         service = self.resolve_table(service) or service
@@ -115,7 +115,7 @@ class MySQL(Vault):
 
     def add_keys(self, service: str, kid_keys: dict[Union[UUID, str], str]) -> int:
         for kid, key in kid_keys.items():
-            if not key or key.count("0") == len(key):
+            if is_null_key(key):
                 raise ValueError("You cannot add a NULL Content Key to a Vault.")
 
         service = self.resolve_table(service) or service

unshackle/vaults/SQLite.py

@@ -5,7 +5,7 @@ from sqlite3 import Connection
 from typing import Iterator, Optional, Union
 from uuid import UUID
 
-from unshackle.core.vault import Vault
+from unshackle.core.vault import Vault, is_null_key
 
 
 class SQLite(Vault):
@@ -58,7 +58,7 @@ class SQLite(Vault):
             cursor.close()
 
     def add_key(self, service: str, kid: Union[UUID, str], key: str) -> bool:
-        if not key or key.count("0") == len(key):
+        if is_null_key(key):
             raise ValueError("You cannot add a NULL Content Key to a Vault.")
 
         service = self.resolve_table(service) or service
@@ -95,7 +95,7 @@ class SQLite(Vault):
 
     def add_keys(self, service: str, kid_keys: dict[Union[UUID, str], str]) -> int:
         for kid, key in kid_keys.items():
-            if not key or key.count("0") == len(key):
+            if is_null_key(key):
                 raise ValueError("You cannot add a NULL Content Key to a Vault.")
 
         service = self.resolve_table(service) or service