diff --git a/docs/reference/configuration/vaults.md b/docs/reference/configuration/vaults.md index de344541..2e478727 100644 --- a/docs/reference/configuration/vaults.md +++ b/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, diff --git a/tests/commands/test_kv_vault_path.py b/tests/commands/test_kv_vault_path.py index 4a7e180a..0e9abaf1 100644 --- a/tests/commands/test_kv_vault_path.py +++ b/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 diff --git a/tests/conftest.py b/tests/conftest.py index 5709f220..c30deac5 100644 --- a/tests/conftest.py +++ b/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.""" diff --git a/tests/core/test_manifest_snapshot.py b/tests/core/test_manifest_snapshot.py index b8ec627d..5473f062 100644 --- a/tests/core/test_manifest_snapshot.py +++ b/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"] diff --git a/tests/core/test_vault_table_case.py b/tests/core/test_vault_table_case.py index 3f3548fc..acb79570 100644 --- a/tests/core/test_vault_table_case.py +++ b/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.""" diff --git a/tests/downloaders/test_multiprocess.py b/tests/downloaders/test_multiprocess.py index 3c56c00c..9918c6ed 100644 --- a/tests/downloaders/test_multiprocess.py +++ b/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 diff --git a/tests/downloaders/test_speed_limit.py b/tests/downloaders/test_speed_limit.py index d468bb87..2a71b941 100644 --- a/tests/downloaders/test_speed_limit.py +++ b/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(): diff --git a/tests/remote/unit/test_serve_reload.py b/tests/remote/unit/test_serve_reload.py index af6fcf1c..c0fcd1db 100644 --- a/tests/remote/unit/test_serve_reload.py +++ b/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"]}}}) diff --git a/tests/remote/unit/test_worker_speed_limit.py b/tests/remote/unit/test_worker_speed_limit.py index d57efd28..7c81ddba 100644 --- a/tests/remote/unit/test_worker_speed_limit.py +++ b/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) diff --git a/unshackle/commands/kv.py b/unshackle/commands/kv.py index a460b9c3..203809f0 100644 --- a/unshackle/commands/kv.py +++ b/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[0-9a-fA-F]{32}):(?P[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[0-9a-fA-F]{32}):(?P[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 diff --git a/unshackle/core/cdm/custom_remote_cdm.py b/unshackle/core/cdm/custom_remote_cdm.py index 2946f9cb..062ffa85 100644 --- a/unshackle/core/cdm/custom_remote_cdm.py +++ b/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 diff --git a/unshackle/core/cdm/decrypt_labs_remote_cdm.py b/unshackle/core/cdm/decrypt_labs_remote_cdm.py index 39268554..978f7b15 100644 --- a/unshackle/core/cdm/decrypt_labs_remote_cdm.py +++ b/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 diff --git a/unshackle/core/drm/__init__.py b/unshackle/core/drm/__init__.py index 8196cd49..22bcc3cd 100644 --- a/unshackle/core/drm/__init__.py +++ b/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) diff --git a/unshackle/core/manifests/snapshot.py b/unshackle/core/manifests/snapshot.py index e667ec0c..dca65781 100644 --- a/unshackle/core/manifests/snapshot.py +++ b/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 diff --git a/unshackle/core/vault.py b/unshackle/core/vault.py index 83b8b5bd..cda12c83 100644 --- a/unshackle/core/vault.py +++ b/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. diff --git a/unshackle/core/vaults.py b/unshackle/core/vaults.py index f60c432f..017388e9 100644 --- a/unshackle/core/vaults.py +++ b/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)) diff --git a/unshackle/vaults/HTTP.py b/unshackle/vaults/HTTP.py index b6c98a09..8059204d 100644 --- a/unshackle/vaults/HTTP.py +++ b/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": diff --git a/unshackle/vaults/MySQL.py b/unshackle/vaults/MySQL.py index 2a2123c2..94384503 100644 --- a/unshackle/vaults/MySQL.py +++ b/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 diff --git a/unshackle/vaults/SQLite.py b/unshackle/vaults/SQLite.py index b8bdac54..510a2421 100644 --- a/unshackle/vaults/SQLite.py +++ b/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