diff --git a/tests/core/test_vault_bad_keys.py b/tests/core/test_vault_bad_keys.py index 573eefb7..1359af4d 100644 --- a/tests/core/test_vault_bad_keys.py +++ b/tests/core/test_vault_bad_keys.py @@ -265,7 +265,10 @@ def test_retry_takes_the_next_parallel_answer_without_a_new_request(tmp_path: Pa assert fast.calls == 1 and slow.calls == 1 -def test_cdm_returning_the_flagged_key_clears_the_flag(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: +@pytest.mark.parametrize("replace", [False, True]) +def test_cdm_returning_the_flagged_key_clears_the_flag( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, replace: bool +) -> None: """A decode failure on a key the CDM then confirms is a false positive: unflag, keep the file.""" local = SQLite("local", tmp_path / "local.db") poisoned = Remote("poisoned", tmp_path / "poisoned.db") @@ -280,13 +283,67 @@ def test_cdm_returning_the_flagged_key_clears_the_flag(tmp_path: Path, monkeypat def licence(drm: FakeDRM, track_kid: UUID) -> None: assert vaults.get_key(KID) == (None, None) - drm.content_keys[KID] = BAD # the CDM answers with the same key + if replace: + drm.content_keys = {KID: BAD} # a CDM licence builds a new dict + else: + drm.content_keys[KID] = BAD # the CDM answers with the same key cmd.decrypt_verified(drm, path, licence, KID) assert path.read_bytes() == b"plain:" + BAD.encode() assert not local.is_bad_key(KID, BAD) +def test_a_replacement_key_in_a_new_dict_is_checked_and_stored(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + local = SQLite("local", tmp_path / "local.db") + poisoned = Remote("poisoned", tmp_path / "poisoned.db") + poisoned.add_key("SVC", KID, BAD) + vaults = Vaults("SVC") + vaults.vaults = [local, poisoned] + assert vaults.get_key(KID) == (BAD, poisoned) + cmd = make_cmd(monkeypatch, vaults) + path = tmp_path / "video.mp4" + path.write_bytes(b"cipher") + drm = FakeDRM({KID: BAD}) + other = Remote("other", tmp_path / "other.db") + other.add_key("SVC", KID, GOOD) + + def licence(drm: FakeDRM, track_kid: UUID) -> None: + vaults.vaults.append(other) + drm.content_keys = {KID: vaults.get_key(KID)[0]} + + cmd.decrypt_verified(drm, path, licence, KID) + cmd.wait_vault_writes() + assert path.read_bytes() == b"plain:" + GOOD.encode() + assert local.get_key(KID, "SVC") == GOOD + + +@pytest.mark.parametrize("replace", [True, False]) +def test_a_cdm_that_confirms_a_key_flagged_on_entry_clears_the_flag( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, replace: bool +) -> None: + """Another track flagged the pair, then this track's CDM licence answers with the same key.""" + local = SQLite("local", tmp_path / "local.db") + vaults = Vaults("SVC") + vaults.vaults = [local] + cmd = make_cmd(monkeypatch, vaults) + vaults.flag_bad_key(KID, BAD) + path = tmp_path / "video.mp4" + path.write_bytes(b"cipher") + drm = FakeDRM({KID: BAD}) + + def licence(drm: FakeDRM, track_kid: UUID) -> None: + if replace: + drm.content_keys = {KID: BAD} + else: + drm.content_keys[KID] = BAD + + cmd.decrypt_verified(drm, path, licence, KID) + cmd.wait_vault_writes() + assert path.read_bytes() == b"plain:" + BAD.encode() + assert not local.is_bad_key(KID, BAD) + assert local.get_key(KID, "SVC") == BAD + + def test_a_licence_key_that_replaces_a_bad_local_pair_is_stored( tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: diff --git a/unshackle/commands/dl.py b/unshackle/commands/dl.py index 04c179e3..37f49716 100644 --- a/unshackle/commands/dl.py +++ b/unshackle/commands/dl.py @@ -4506,7 +4506,7 @@ class dl: With ``decrypt=False`` the downloader already decrypted the segments in place, so a failure can only flag and raise; the next run skips the flagged pair. """ - keys = getattr(drm, "content_keys", None) + keys: Any = getattr(drm, "content_keys", None) if keys is None: if decrypt: drm.decrypt(path) @@ -4522,25 +4522,39 @@ class dl: self.LICENSE_KEY_CACHE.pop(kid) return key - def relicense(dropped: set[UUID]) -> None: - """Licence again, then ask for each dropped KID that is still without a content key. + def relicense(bad: dict[UUID, str], kids: dict[UUID, Vault]) -> None: + """Licence again for each dropped KID, then settle each dropped pair against the new key. prepare_drm licenses only for a KID it is told to need, so without the second call a - track that still holds other content keys would decrypt without the dropped one. + track that still holds other content keys would decrypt without the dropped one. A + licence can replace the DRM's dict, so ``keys`` is read from the DRM again afterwards. + A licence key equal to the dropped one proves the decode check wrong and clears its flag. """ + nonlocal keys if not licence: return licence(drm, track_kid=track_kid) - for kid in sorted(dropped): - # prepare_drm can replace the dict, so read it from the DRM each time + for kid in sorted(bad): if kid != track_kid and kid not in getattr(drm, "content_keys", {}): licence(drm, track_kid=kid) + keys = getattr(drm, "content_keys", keys) + for kid, key in bad.items(): + new = keys.get(kid) + if new is None: + continue + source = self.vaults.sources.get(kid) + if source is None: + if new == key: + self.vaults.unflag_bad_key(kid, key) + self.flush_vault_writes([partial(self.vaults.add_key, kid, key)]) + elif not self.vaults.held_locally(kid, new): + self.flush_vault_writes([partial(self.vaults.add_key, kid, new)]) + elif new == key and kid in kids and source[1] is kids[kid]: + raise ValueError(f"{key} from {kids[kid].name} was bad and no other source has the key") known_bad = flagged_kids() if known_bad and licence: - for kid in known_bad: - drop(kid) - relicense(set(known_bad)) + relicense({kid: drop(kid) for kid in known_bad}, {}) def vault_kids() -> dict[UUID, Vault]: """The KIDs whose key on the track came from a vault, with that vault.""" @@ -4625,20 +4639,7 @@ class dl: raise ValueError("The content key from the vault did not decrypt the track; run again") path.unlink() os.link(backup, path) - relicense(set(bad)) - for kid, key in bad.items(): - new = keys.get(kid) - if new is None: - continue - source = self.vaults.sources.get(kid) - if source is None: - if new == key: - self.vaults.unflag_bad_key(kid, key) - self.flush_vault_writes([partial(self.vaults.add_key, kid, key)]) - elif not self.vaults.held_locally(kid, new): - self.flush_vault_writes([partial(self.vaults.add_key, kid, new)]) - elif new == key and kid in kids and source[1] is kids[kid]: - raise ValueError(f"{key} from {kids[kid].name} was bad and no other source has the key") + relicense(bad, kids) raise ValueError("No vault or CDM produced a content key that decrypts the track") finally: backup.unlink(missing_ok=True)