fix(dl): read the drm's content keys again after a relicence replaces them

A licence that builds a new key dict left decrypt_verified on the old one, so a falsely flagged key stayed flagged and a replacement key went unchecked. Both relicence passes now share one step that clears the flag.

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:

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)