unshackle/tests/core/test_vault_bad_keys.py

import logging
import threading
import time
from pathlib import Path
from types import SimpleNamespace
from typing import Optional, Union
from uuid import UUID

import pytest

from tests.remote.key_source import remote_keys
from unshackle.commands import dl as dl_module
from unshackle.commands.dl import dl
from unshackle.core.vault import RateLimited, vault_session
from unshackle.core.vaults import Vaults
from unshackle.vaults.API import API, Exceptions
from unshackle.vaults.HTTP import HTTP
from unshackle.vaults.SQLite import SQLite

KID = UUID("11111111-2222-3333-4444-555555555555")
SIBLING = UUID("99999999-2222-3333-4444-555555555555")
BAD = "aa" * 16
GOOD = "bb" * 16


class Remote(SQLite):
    """A second SQLite file standing in for a remote vault."""

    local = False


def test_flagged_pair_is_skipped_and_named(tmp_path: Path) -> 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)
    assert vaults.sources[KID] == (BAD, poisoned)
    local.add_key("SVC", KID, BAD)  # the cross-vault copy that happened before verification

    vaults.flag_bad_key(KID, BAD)

    # both the local copy and every later remote hit of that pair are gone
    assert local.get_key(KID, "SVC") is None
    assert vaults.get_key(KID) == (None, None)
    assert local.is_bad_key(KID, BAD)
    assert "bad_keys" not in list(local.get_services())
    row = local.conn_factory.get().execute("SELECT service, source FROM bad_keys").fetchone()
    assert row == ("SVC", "poisoned")

    # another vault with a different key for the same KID still answers
    other = Remote("other", tmp_path / "other.db")
    other.add_key("SVC", KID, GOOD)
    vaults.vaults.append(other)
    assert vaults.get_key(KID) == (GOOD, other)

    # a flagged row that comes back (kv copy from an old export) is invisible to the lookup
    local.conn_factory.get().execute("INSERT OR IGNORE INTO `SVC` (kid, key_) VALUES (?, ?)", (KID.hex, BAD))
    local.conn_factory.get().commit()
    assert local.get_key(KID, "SVC") is None

    vaults.unflag_bad_key(KID, BAD)
    assert not local.is_bad_key(KID, BAD)
    assert local.get_key(KID, "SVC") == BAD


class FakeDRM:
    def __init__(self, keys: dict[UUID, str]):
        self.content_keys = dict(keys)
        self.decrypted_with: list[dict[UUID, str]] = []

    def decrypt(self, path: Path) -> None:
        self.decrypted_with.append(dict(self.content_keys))
        # like the real decrypters: write a new file, then replace the input
        out = path.with_stem(path.stem + "_decrypted")
        out.write_bytes(b"plain:" + self.content_keys[KID].encode())
        path.unlink()
        out.rename(path)


def make_cmd(monkeypatch: pytest.MonkeyPatch, vaults: Vaults) -> dl:
    monkeypatch.setattr(dl_module.binaries, "FFMPEG", "ffmpeg")
    monkeypatch.setattr(
        dl_module,
        "ffmpeg_decodes",
        lambda path, start=None, seconds=3, video=None: path.read_bytes() == b"plain:" + GOOD.encode(),
    )
    cmd = dl.__new__(dl)
    cmd.log = SimpleNamespace(warning=print, info=print, debug=print)
    cmd.vaults = vaults
    cmd.vault_cache_tally = None
    dl.LICENSE_KEY_CACHE.clear()
    dl.LICENSE_KEY_CACHE[KID] = BAD
    return cmd


def test_decrypt_verified_flags_and_retries(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:
        assert KID not in drm.content_keys
        assert vaults.get_key(KID) == (None, None)  # the flagged pair is skipped on the retry
        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 drm.decrypted_with == [{KID: BAD}, {KID: GOOD}]
    assert local.is_bad_key(KID, BAD)
    assert local.get_key(KID, "SVC") == GOOD  # the key that proved itself on the retry is kept locally
    assert KID not in dl.LICENSE_KEY_CACHE
    assert not (tmp_path / "video.mp4.enc").exists()

    # a sibling track that still carries the poisoned key is checked too and takes the cached good key
    sibling = tmp_path / "audio.mp4"
    sibling.write_bytes(b"cipher")
    drm2 = FakeDRM({KID: BAD})
    cmd.decrypt_verified(drm2, sibling, lambda drm, track_kid: drm.content_keys.update({KID: GOOD}), KID)
    assert sibling.read_bytes() == b"plain:" + GOOD.encode()


def test_no_local_vault_raises_instead_of_reusing_the_poison(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
    poisoned = Remote("poisoned", tmp_path / "poisoned.db")
    poisoned.add_key("SVC", KID, BAD)
    vaults = Vaults("SVC")
    vaults.vaults = [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})

    def licence(drm: FakeDRM, track_kid: UUID) -> None:
        key, _ = vaults.get_key(KID)  # nothing holds the flag, so the same pair comes back
        drm.content_keys[KID] = key

    with pytest.raises(ValueError, match="no other source"):
        cmd.decrypt_verified(drm, path, licence, KID)
    assert path.read_bytes() == b"cipher"  # ciphertext restored, nothing written with the poison
    assert not (tmp_path / "video.mp4.enc").exists()


def test_segment_decrypted_track_flags_and_raises(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]
    vaults.get_key(KID)
    cmd = make_cmd(monkeypatch, vaults)
    path = tmp_path / "video.mp4"
    path.write_bytes(b"plain:" + BAD.encode())
    with pytest.raises(ValueError, match="run again"):
        cmd.decrypt_verified(FakeDRM({KID: BAD}), path, lambda *a, **k: None, KID, decrypt=False)
    assert local.is_bad_key(KID, BAD)


def test_downloaders_route_through_the_hook(monkeypatch: pytest.MonkeyPatch) -> None:
    """Every real decrypt site calls decrypt_track, so the dl hook sees it."""
    import inspect

    from unshackle.core.manifests import dash, hls, ism
    from unshackle.core.tracks import track

    for module in (dash, hls, ism, track):
        src = inspect.getsource(module)
        assert "decrypt_track(" in src, module.__name__
        assert "drm.decrypt(save_path)" not in src and "drm.decrypt(merged_path)" not in src, module.__name__


def test_every_poisoned_vault_is_burned_then_the_cdm_answers(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
    local = SQLite("local", tmp_path / "local.db")
    poisoned = [Remote(f"poisoned{i}", tmp_path / f"p{i}.db") for i in range(3)]
    for i, vault in enumerate(poisoned):
        vault.add_key("SVC", KID, f"{i + 1:02x}" * 16)
    vaults = Vaults("SVC")
    vaults.vaults = [local, *poisoned]
    cmd = make_cmd(monkeypatch, vaults)
    path = tmp_path / "video.mp4"
    path.write_bytes(b"cipher")
    drm = FakeDRM({KID: vaults.get_key(KID)[0]})
    cdm_calls = 0

    def licence(drm: FakeDRM, track_kid: UUID) -> None:
        nonlocal cdm_calls
        key, _ = vaults.get_key(KID)
        if key is None:
            cdm_calls += 1
            key = GOOD
        drm.content_keys[KID] = key

    cmd.decrypt_verified(drm, path, licence, KID)
    assert path.read_bytes() == b"plain:" + GOOD.encode()
    assert cdm_calls == 1
    assert len(drm.decrypted_with) == 4  # three vault keys burned, then the CDM key
    assert all(local.is_bad_key(KID, f"{i + 1:02x}" * 16) for i in range(3))


def test_kv_copy_never_moves_the_flag_table(tmp_path: Path) -> None:
    import logging

    from unshackle.commands.kv import copy_service_data

    src = SQLite("src", tmp_path / "src.db")
    dst = SQLite("dst", tmp_path / "dst.db")
    src.add_key("SVC", KID, GOOD)
    src.flag_bad_key("SVC", KID, BAD, "poisoned")
    assert list(src.get_services()) == ["SVC"]
    assert copy_service_data(dst, src, "bad_keys", logging.getLogger("t")) == 0
    assert list(dst.get_services()) == []
    # a flagged pair in the source never reaches a destination that knows it is bad
    kid2 = UUID("22222222-2222-3333-4444-555555555555")
    src.add_key("SVC", kid2, BAD)
    dst.flag_bad_key("SVC", kid2, BAD, "poisoned")
    assert copy_service_data(dst, src, "SVC", logging.getLogger("t")) == 1
    assert list(dst.get_keys("SVC")) == [(KID.hex, GOOD)]


def test_retry_takes_the_next_parallel_answer_without_a_new_request(tmp_path: Path) -> None:
    from tests.core.test_vaults_lookup import FakeVault

    started, release = threading.Event(), threading.Event()

    class Fast(FakeVault):
        def get_key(self, kid: Union[UUID, str], service: str) -> Optional[str]:
            # answer only once the slow query runs, or the pool shutdown cancels it before it starts
            started.wait(5)
            return super().get_key(kid, service)

    class Slow(FakeVault):
        def get_key(self, kid: Union[UUID, str], service: str) -> Optional[str]:
            started.set()
            release.wait(5)
            return super().get_key(kid, service)

    local = SQLite("local", tmp_path / "local.db")
    fast, slow = Fast("fast", BAD), Slow("slow", GOOD)
    vaults = Vaults("SVC")
    vaults.vaults = [local, fast, slow]

    assert vaults.get_key(KID) == (BAD, fast)
    vaults.flag_bad_key(KID, BAD)
    release.set()  # the slow answer lands in candidates after the first call returned
    deadline = time.monotonic() + 5
    while len(vaults.candidates.get(KID, [])) < 2 and time.monotonic() < deadline:
        time.sleep(0.01)
    assert vaults.get_key(KID) == (GOOD, slow)
    assert fast.calls == 1 and slow.calls == 1


@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")
    poisoned.add_key("SVC", KID, BAD)
    vaults = Vaults("SVC")
    vaults.vaults = [local, poisoned]
    cmd = make_cmd(monkeypatch, vaults)
    monkeypatch.setattr(dl_module, "ffmpeg_decodes", lambda path, start=None, seconds=3, video=None: False)
    path = tmp_path / "video.mp4"
    path.write_bytes(b"cipher")
    drm = FakeDRM({KID: vaults.get_key(KID)[0]})

    def licence(drm: FakeDRM, track_kid: UUID) -> None:
        assert vaults.get_key(KID) == (None, None)
        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:
    """The local vault held the bad pair when the licence key arrived, so it skipped the KID; store it now."""
    local = SQLite("local", tmp_path / "local.db")
    local.add_key("SVC", KID, BAD)
    vaults = Vaults("SVC")
    vaults.vaults = [local]
    cmd = make_cmd(monkeypatch, vaults)
    drm = FakeDRM({KID: vaults.get_key(KID)[0]})
    path = tmp_path / "video.mp4"
    path.write_bytes(b"cipher")
    assert local.add_keys("SVC", {KID: GOOD}) == 0

    cmd.decrypt_verified(drm, path, lambda drm, track_kid: drm.content_keys.update({KID: GOOD}), KID)
    cmd.wait_vault_writes()
    assert local.is_bad_key(KID, BAD)
    assert local.get_key(KID, "SVC") == GOOD


def test_stale_backup_from_a_killed_run_is_not_restored(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
    poisoned = Remote("poisoned", tmp_path / "poisoned.db")
    poisoned.add_key("SVC", KID, BAD)
    vaults = Vaults("SVC")
    vaults.vaults = [poisoned]
    vaults.get_key(KID)
    cmd = make_cmd(monkeypatch, vaults)
    path = tmp_path / "video.mp4"
    path.write_bytes(b"cipher")
    (tmp_path / "video.mp4.enc").write_bytes(b"stale")

    def licence(drm: FakeDRM, track_kid: UUID) -> None:
        drm.content_keys[KID] = vaults.get_key(KID)[0]

    with pytest.raises(ValueError, match="no other source"):
        cmd.decrypt_verified(FakeDRM({KID: BAD}), path, licence, KID)
    assert path.read_bytes() == b"cipher"


def test_drm_without_content_keys_still_decrypts(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
    """HLS AES-128 ClearKey has no KID:KEY map; the hook must pass it straight through."""
    cmd = make_cmd(monkeypatch, Vaults("SVC"))
    path = tmp_path / "video.ts"
    path.write_bytes(b"cipher")
    calls: list[Path] = []
    cmd.decrypt_verified(SimpleNamespace(decrypt=calls.append), path)
    assert calls == [path]


def test_flag_on_a_read_only_vault_keeps_its_rows(tmp_path: Path) -> None:
    ro = SQLite("ro", tmp_path / "ro.db", no_push=True)
    ro.add_key("SVC", KID, BAD)
    ro.flag_bad_key("SVC", KID, BAD, "poisoned")
    assert ro.is_bad_key(KID, BAD)
    assert ro.get_key(KID, "SVC") is None  # the lookup still skips the flagged pair
    assert list(ro.get_keys("SVC")) == [(KID.hex, BAD)]  # but the row itself is untouched


def test_flag_moves_the_row_and_blocks_it_coming_back(tmp_path: Path) -> None:
    local = SQLite("local", tmp_path / "local.db")
    local.add_key("SVC", KID, BAD)
    local.flag_bad_key("SVC", KID, BAD, "poisoned")
    assert local.is_bad_key(KID, BAD)
    assert list(local.get_keys("SVC")) == []
    assert local.add_key("SVC", KID, BAD) is False
    assert local.add_keys("SVC", {KID: BAD}) == 0
    assert list(local.get_keys("SVC")) == []
    assert local.add_key("SVC", KID, GOOD)  # the KID itself is not blocked, only the pair


def test_a_licence_key_clears_a_flag_from_an_earlier_run(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
    """A false flag must not outlive the run: the next licence returns the pair and it is stored again."""
    local = SQLite("local", tmp_path / "local.db")
    local.flag_bad_key("SVC", KID, GOOD, "poisoned")
    vaults = Vaults("SVC")
    vaults.vaults = [local]
    cmd = make_cmd(monkeypatch, vaults)
    cmd.vault_cache_tally = None
    cmd.cache_keys_to_vaults({KID: GOOD})
    assert not local.is_bad_key(KID, GOOD)
    assert local.get_key(KID, "SVC") == GOOD


def test_a_licence_with_no_new_keys_sends_no_vault_write(monkeypatch: pytest.MonkeyPatch) -> None:
    """A PSSH KID the licence never answers forces a licence on every run; the vaults must not get an empty push."""
    vault = SimpleNamespace(no_push=False, add_keys=lambda *_: pytest.fail("empty push reached a vault"))
    vaults = Vaults("SVC")
    vaults.vaults = [vault]  # type: ignore[list-item]
    cmd = make_cmd(monkeypatch, vaults)
    cmd.cache_keys_to_vaults({})
    assert cmd.vault_cache_tally is None


def test_the_cached_key_count_is_unique_and_skipped_without_vaults(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    """Two tracks licensed in parallel can return the same key; --cdm-only loads no vaults at all."""
    other = UUID(int=9)
    vaults = Vaults("SVC")
    vaults.vaults = [SQLite("local", tmp_path / "local.db")]
    cmd = make_cmd(monkeypatch, vaults)
    cmd.cache_keys_to_vaults({KID: GOOD})
    cmd.cache_keys_to_vaults({KID: GOOD, other: GOOD})
    assert cmd.vault_cache_tally == ({KID, other}, 1)

    cmd = make_cmd(monkeypatch, Vaults("SVC"))
    cmd.cache_keys_to_vaults({KID: GOOD})
    assert cmd.vault_cache_tally is None


def test_bulk_add_refuses_a_flagged_pair_in_any_case_or_kid_form(tmp_path: Path) -> None:
    local = SQLite("local", tmp_path / "local.db")
    local.flag_bad_key("SVC", KID, BAD, "poisoned")
    for kid, key in ((KID.hex.upper(), BAD), (str(KID), BAD.upper()), (KID, BAD)):
        assert local.add_keys("SVC", {kid: key}) == 0
    assert list(local.get_keys("SVC")) == []


def test_only_the_kid_whose_window_fails_is_flagged(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
    """A lead under one KID and the rest under another: the lead's key is right and stays."""
    lead = UUID(int=7)
    local = SQLite("local", tmp_path / "local.db")
    poisoned = Remote("poisoned", tmp_path / "poisoned.db")
    poisoned.add_key("SVC", lead, BAD)  # the same wrong key under both KIDs
    poisoned.add_key("SVC", KID, BAD)
    other = Remote("other", tmp_path / "other.db")
    other.add_key("SVC", KID, GOOD)
    vaults = Vaults("SVC")
    vaults.vaults = [local, poisoned]
    cmd = make_cmd(monkeypatch, vaults)
    drm = FakeDRM({lead: vaults.get_key(lead)[0], KID: vaults.get_key(KID)[0]})
    right = {lead: BAD, KID: GOOD}
    monkeypatch.setattr(
        dl_module.verify,
        "kid_windows",
        lambda path: dl_module.verify.KidMap({lead: [(0.0, 2.0)], KID: [(120.0, 122.0)]}, True),
    )
    monkeypatch.setattr(
        dl_module,
        "ffmpeg_decodes",
        lambda path, start=None, seconds=3, video=None: (
            drm.decrypted_with[-1][lead if start == 0.0 else KID] == right[lead if start == 0.0 else KID]
        ),
    )
    path = tmp_path / "video.mp4"
    path.write_bytes(b"cipher")

    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 drm.content_keys == right
    assert local.is_bad_key(KID, BAD)
    assert not local.is_bad_key(lead, BAD)


def test_a_kid_absent_from_the_map_is_judged_by_the_fallback(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
    """A map that holds an all-zero tenc KID is not complete, so a KID it does not know
    shares the whole-file verdict, and a wrong key there is flagged."""
    lead = UUID(int=7)
    local = SQLite("local", tmp_path / "local.db")
    poisoned = Remote("poisoned", tmp_path / "poisoned.db")
    poisoned.add_key("SVC", lead, GOOD)
    poisoned.add_key("SVC", KID, BAD)
    other = Remote("other", tmp_path / "other.db")
    other.add_key("SVC", KID, GOOD)
    vaults = Vaults("SVC")
    vaults.vaults = [local, poisoned]
    cmd = make_cmd(monkeypatch, vaults)
    drm = FakeDRM({lead: vaults.get_key(lead)[0], KID: vaults.get_key(KID)[0]})
    calls: list[tuple] = []

    def decodes(path: Path, start: object = None, seconds: float = 3, video: object = None) -> bool:
        calls.append((start, seconds, video))
        return drm.decrypted_with[-1][lead if start is not None else KID] == GOOD

    monkeypatch.setattr(
        dl_module.verify,
        "kid_windows",
        lambda path: dl_module.verify.KidMap({lead: [(4.0, 6.0)], UUID(int=0): [(8.0, 10.0)]}, True),
    )
    monkeypatch.setattr(dl_module, "ffmpeg_decodes", decodes)
    path = tmp_path / "video.mp4"
    path.write_bytes(b"cipher")

    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 drm.content_keys == {lead: GOOD, KID: GOOD}
    assert local.is_bad_key(KID, BAD)
    assert not local.is_bad_key(lead, GOOD)
    assert (4.0, 2.0, True) in calls  # the window is capped at the fragment's end
    assert (None, 3, True) in calls  # the fallback knows the track is video


def test_a_wrong_mapped_key_does_not_flag_a_kid_the_file_does_not_carry(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    """The audio KID is in the licence and the vault but not in this video file: the whole-file
    check fails on the wrong video key's noise, and that verdict must not reach the audio KID."""
    audio = UUID(int=7)
    local = SQLite("local", tmp_path / "local.db")
    poisoned = Remote("poisoned", tmp_path / "poisoned.db")
    poisoned.add_key("SVC", audio, GOOD)
    poisoned.add_key("SVC", KID, BAD)
    other = Remote("other", tmp_path / "other.db")
    other.add_key("SVC", KID, GOOD)
    vaults = Vaults("SVC")
    vaults.vaults = [local, poisoned]
    cmd = make_cmd(monkeypatch, vaults)
    drm = FakeDRM({audio: vaults.get_key(audio)[0], KID: vaults.get_key(KID)[0]})
    monkeypatch.setattr(
        dl_module.verify, "kid_windows", lambda path: dl_module.verify.KidMap({KID: [(4.0, 6.0)]}, True)
    )
    monkeypatch.setattr(
        dl_module,
        "ffmpeg_decodes",
        lambda path, start=None, seconds=3, video=None: drm.decrypted_with[-1][KID] == GOOD,
    )
    path = tmp_path / "video.mp4"
    path.write_bytes(b"cipher")

    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 drm.content_keys == {audio: GOOD, KID: GOOD}
    assert local.is_bad_key(KID, BAD)
    assert not local.is_bad_key(audio, GOOD)
    assert local.get_key(audio, "SVC") is None


def test_a_kid_the_file_does_not_carry_is_not_judged_by_it(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
    """The CDM licensed the video KID and a vault served the audio KID. An FFmpeg that fails
    every decode of the video file must not flag the audio KID, nor copy it to another vault."""
    audio = UUID(int=7)
    local = SQLite("local", tmp_path / "local.db")
    remote = Remote("remote", tmp_path / "remote.db")
    remote.add_key("SVC", audio, GOOD)
    vaults = Vaults("SVC")
    vaults.vaults = [local, remote]
    cmd = make_cmd(monkeypatch, vaults)
    drm = FakeDRM({audio: vaults.get_key(audio)[0], KID: GOOD})
    monkeypatch.setattr(
        dl_module.verify, "kid_windows", lambda path: dl_module.verify.KidMap({KID: [(4.0, 6.0)]}, True)
    )
    monkeypatch.setattr(dl_module, "ffmpeg_decodes", lambda path, start=None, seconds=3, video=None: False)
    path = tmp_path / "video.mp4"
    path.write_bytes(b"cipher")
    licences: list[UUID] = []

    cmd.decrypt_verified(drm, path, lambda drm, track_kid: licences.append(track_kid), KID)
    cmd.wait_vault_writes()
    assert drm.content_keys == {audio: GOOD, KID: GOOD}
    assert not licences
    assert not local.is_bad_key(audio, GOOD)
    assert local.get_key(audio, "SVC") is None


@pytest.mark.parametrize(
    ("served", "decodes", "told"),
    [
        (GOOD, False, False),  # broken FFmpeg, the CDM returns the same key
        (BAD, True, True),  # the CDM returns another key
    ],
)
def test_a_remote_vault_is_told_only_when_another_key_replaces_its_pair(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch, served: str, decodes: bool, told: bool
) -> None:
    remote = Remote("remote", tmp_path / "remote.db")
    remote.add_key("SVC", KID, served)
    vaults = Vaults("SVC")
    vaults.vaults = [SQLite("local", tmp_path / "local.db"), remote]
    cmd = make_cmd(monkeypatch, vaults)
    if not decodes:
        monkeypatch.setattr(dl_module, "ffmpeg_decodes", lambda path, start=None, seconds=3, video=None: False)
    drm = FakeDRM({KID: vaults.get_key(KID)[0]})
    path = tmp_path / "video.mp4"
    path.write_bytes(b"cipher")

    def licence(drm: FakeDRM, track_kid: UUID) -> None:
        assert not remote.is_bad_key(KID, served)
        drm.content_keys[KID] = GOOD

    cmd.decrypt_verified(drm, path, licence, KID)
    assert drm.content_keys == {KID: GOOD}
    assert remote.is_bad_key(KID, served) is told


def test_the_server_vault_is_told_at_once(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
    """Over --remote the report is what makes the server skip the pair on the next licence."""
    reports: list[tuple[UUID, str]] = []
    server = SimpleNamespace(
        name="server_cdm",
        local=False,
        report_now=True,
        flag_bad_key=lambda service, kid, key, source: reports.append((kid, key)),
    )
    vaults = Vaults("SVC")
    vaults.vaults = [SQLite("local", tmp_path / "local.db")]
    vaults.sources[KID] = (BAD, server)
    cmd = make_cmd(monkeypatch, vaults)
    drm = FakeDRM({KID: BAD})
    path = tmp_path / "video.mp4"
    path.write_bytes(b"cipher")

    def licence(drm: FakeDRM, track_kid: UUID) -> None:
        assert reports == [(KID, BAD)]
        drm.content_keys[KID] = GOOD

    cmd.decrypt_verified(drm, path, licence, KID)
    assert drm.content_keys == {KID: GOOD}


def test_only_checked_kids_are_copied_to_other_vaults(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
    """Only the mapped KID is proven and copied; the KID this file does not carry waits for its own file."""
    lead = UUID(int=7)
    local = SQLite("local", tmp_path / "local.db")
    remote = Remote("remote", tmp_path / "remote.db")
    remote.add_key("SVC", lead, GOOD)
    remote.add_key("SVC", KID, GOOD)
    vaults = Vaults("SVC")
    vaults.vaults = [local, remote]
    cmd = make_cmd(monkeypatch, vaults)
    drm = FakeDRM({lead: vaults.get_key(lead)[0], KID: vaults.get_key(KID)[0]})
    starts: list = []

    def decodes(path: Path, start: object = None, seconds: float = 3, video: object = None) -> bool:
        starts.append(start)
        return True

    monkeypatch.setattr(
        dl_module.verify, "kid_windows", lambda path: dl_module.verify.KidMap({lead: [(4.0, 6.0)]}, True)
    )
    monkeypatch.setattr(dl_module, "ffmpeg_decodes", decodes)
    path = tmp_path / "video.mp4"
    path.write_bytes(b"cipher")

    cmd.decrypt_verified(drm, path, None, KID)
    cmd.wait_vault_writes()
    assert None not in starts
    assert local.get_key(lead, "SVC") == GOOD
    assert local.get_key(KID, "SVC") is None


def test_a_clear_file_judges_and_copies_no_vault_key(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
    local = SQLite("local", tmp_path / "local.db")
    remote = Remote("remote", tmp_path / "remote.db")
    remote.add_key("SVC", KID, BAD)
    vaults = Vaults("SVC")
    vaults.vaults = [local, remote]
    cmd = make_cmd(monkeypatch, vaults)
    drm = FakeDRM({KID: vaults.get_key(KID)[0]})
    checks: list = []

    def decodes(path: Path, start: object = None, seconds: float = 3, video: object = None) -> bool:
        checks.append(start)
        return True

    monkeypatch.setattr(dl_module.verify, "kid_windows", lambda path: dl_module.verify.KidMap({}, False))
    monkeypatch.setattr(dl_module, "ffmpeg_decodes", decodes)
    path = tmp_path / "audio.mp4"
    path.write_bytes(b"clear")

    cmd.decrypt_verified(drm, path, None, KID)
    cmd.wait_vault_writes()
    assert checks == []
    assert local.get_key(KID, "SVC") is None
    assert not local.is_bad_key(KID, BAD)


def test_flagged_pair_warns_once_per_vault(tmp_path: Path, caplog: pytest.LogCaptureFixture) -> None:
    local = SQLite("local", tmp_path / "local.db")
    poisoned = Remote("poisoned", tmp_path / "poisoned.db")
    poisoned.add_key("SVC", KID, BAD)
    local.flag_bad_key("SVC", KID, BAD, "poisoned")
    vaults = Vaults("SVC")
    vaults.vaults = [local, poisoned]
    with caplog.at_level(logging.WARNING, logger="unshackle.core.vaults"):
        for _ in range(3):
            assert vaults.query_vault(poisoned, KID) is None
    assert sum("flagged bad, skipping" in r.getMessage() for r in caplog.records) == 1


def test_add_key_refuses_a_flagged_pair_with_a_dashed_kid(tmp_path: Path) -> None:
    local = SQLite("local", tmp_path / "local.db")
    local.flag_bad_key("SVC", KID, BAD, "poisoned")
    assert local.add_key("SVC", str(KID), BAD) is False
    assert local.add_key("SVC", str(KID).upper(), BAD.upper()) is False
    assert list(local.get_keys("SVC")) == []


def test_a_sibling_verdict_during_the_decrypt_still_triggers_the_retry(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    """Two tracks share a KID and decrypt at the same time; the first verdict pops the shared
    source, so the second track must read the flag instead of returning unchecked."""
    local = SQLite("local", tmp_path / "local.db")
    vaults = Vaults("SVC")
    vaults.vaults = [local]
    vaults.sources[KID] = (BAD, local)
    cmd = make_cmd(monkeypatch, vaults)
    dl.LICENSE_KEY_CACHE[KID] = GOOD  # the sibling's retry already licensed the good key

    path = tmp_path / "audio.mp4"
    path.write_bytes(b"cipher")

    class SiblingFlagsMeanwhile(FakeDRM):
        def decrypt(self, path: Path) -> None:
            super().decrypt(path)
            if self.content_keys[KID] == BAD:
                vaults.flag_bad_key(KID, BAD)

    drm = SiblingFlagsMeanwhile({KID: BAD})
    licences: list = []

    def licence(drm: FakeDRM, track_kid: UUID) -> None:
        licences.append(dl.LICENSE_KEY_CACHE.get(KID))
        drm.content_keys[KID] = dl.LICENSE_KEY_CACHE[KID]

    cmd.decrypt_verified(drm, path, licence, KID)

    assert path.read_bytes() == b"plain:" + GOOD.encode()
    assert drm.decrypted_with == [{KID: BAD}, {KID: GOOD}]
    assert licences == [GOOD]  # the good key in the run cache survives the flag of the bad one


def test_api_vault_reports_the_pair_upstream(monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture) -> None:
    api = API("vault", "https://vault.example/", "tok")
    calls: list[tuple[str, str, object]] = []
    ok = SimpleNamespace(status_code=200, raise_for_status=lambda: None, json=lambda: {"code": 0})
    monkeypatch.setattr(api.session, "post", lambda url, **kw: calls.append(("POST", url, kw.get("json"))) or ok)
    monkeypatch.setattr(api.session, "delete", lambda url, **kw: calls.append(("DELETE", url, None)) or ok)

    vaults = Vaults("SVC")
    vaults.vaults = [api]
    vaults.sources[KID] = (BAD, api)
    vaults.flag_bad_key(KID, BAD)
    assert not calls
    vaults.report_replaced({KID: GOOD})
    vaults.unflag_bad_key(KID, BAD)

    assert calls == [
        ("POST", f"https://vault.example/bad-keys/{KID.hex}", {"content_key": BAD, "service": "SVC"}),
        ("DELETE", f"https://vault.example/bad-keys/{KID.hex}/{BAD}", None),
    ]

    def flag_and_replace(source: object) -> None:
        vaults.sources[KID] = (BAD, source)
        vaults.flag_bad_key(KID, BAD)
        vaults.report_replaced({KID: GOOD})

    calls.clear()
    flag_and_replace(SimpleNamespace(name="other", local=False, report_now=False))
    assert calls == []

    monkeypatch.setattr(api.session, "post", lambda url, **kw: (_ for _ in ()).throw(ConnectionError("down")))
    caplog.clear()
    with caplog.at_level(logging.DEBUG, logger="unshackle.core.vaults"):
        flag_and_replace(api)
    assert [r.levelname for r in caplog.records] == ["DEBUG"]

    gone = SimpleNamespace(status_code=404, raise_for_status=lambda: (_ for _ in ()).throw(RuntimeError("404")))
    monkeypatch.setattr(api.session, "post", lambda url, **kw: gone)
    caplog.clear()
    flag_and_replace(api)
    assert not caplog.records

    denied = SimpleNamespace(status_code=200, raise_for_status=lambda: None, json=lambda: {"code": 1, "message": "no"})
    monkeypatch.setattr(api.session, "post", lambda url, **kw: denied)
    monkeypatch.setattr(api.session, "delete", lambda url, **kw: denied)
    flag_and_replace(api)
    vaults.unflag_bad_key(KID, BAD)
    assert not caplog.records


def test_prefer_vault_keys_queries_each_named_kid_once() -> None:
    asked: list[UUID] = []
    cmd = dl.__new__(dl)
    cmd.vaults = SimpleNamespace(get_key=lambda kid: asked.append(kid) or (None, None))
    video, audio = (SimpleNamespace(kids=[KID], content_keys={KID: GOOD, SIBLING: BAD}) for _ in range(2))

    cmd.prefer_vault_keys(
        remote_keys(), SimpleNamespace(tracks=[SimpleNamespace(drm=[video]), SimpleNamespace(drm=[audio])])
    )

    assert asked == [KID]


def test_cache_resolved_keys_skips_a_key_a_local_vault_holds(tmp_path: Path) -> None:
    local = SQLite("local", tmp_path / "local.db")
    local.add_key("SVC", SIBLING, BAD)
    cmd = dl.__new__(dl)
    cmd.vaults = Vaults("SVC")
    cmd.vaults.vaults = [local]
    pushed: list[dict] = []
    cmd.flush_vault_writes = lambda writes: pushed.extend(w.args[0] for w in writes)
    dl.LICENSE_KEY_CACHE.clear()
    drm = SimpleNamespace(kids=[KID], content_keys={KID: GOOD, SIBLING: BAD})

    cmd.cache_resolved_keys(remote_keys(), SimpleNamespace(tracks=[SimpleNamespace(drm=[drm])]))

    assert pushed == [{KID: GOOD}]


def test_api_vault_does_not_shrink_the_batch_on_429(monkeypatch: pytest.MonkeyPatch) -> None:
    api = API("vault", "https://vault.example/", "tok")
    posts: list[str] = []
    limited = SimpleNamespace(status_code=429)
    monkeypatch.setattr(api.session, "post", lambda url, **kw: posts.append(url) or limited)

    with pytest.raises(Exceptions.TooManyRequests):
        api.add_keys("SVC", {KID: GOOD})

    assert len(posts) == 1


def test_http_vault_does_not_fall_back_to_single_inserts_on_429(monkeypatch: pytest.MonkeyPatch) -> None:
    vault = HTTP("vault", "https://vault.example/", password="tok", api_mode="json")
    posts: list[str] = []
    limited = SimpleNamespace(status_code=429, ok=False)
    monkeypatch.setattr(vault.session, "post", lambda url, **kw: posts.append(url) or limited)

    with pytest.raises(RateLimited):
        vault.add_keys("SVC", {KID: GOOD, SIBLING: BAD})

    assert len(posts) == 1


def test_vault_session_retries_429_but_not_a_dead_vault() -> None:
    retry = vault_session(10).get_adapter("https://vault.example").max_retries
    assert retry.status_forcelist == [429]
    assert (retry.connect, retry.read, retry.other) == (0, 0, 0)


def test_an_unreachable_vault_costs_one_attempt_per_run(monkeypatch: pytest.MonkeyPatch) -> None:
    import requests

    def down(url: str, **kw: object) -> None:
        calls.append(url)
        raise requests.ConnectTimeout("no answer")

    calls: list[str] = []
    api = API("api", "https://api.example/", "tok")
    http = HTTP("http", "https://http.example/", password="tok", api_mode="json")
    for vault in (api, http):
        monkeypatch.setattr(vault.session, "get", down)
        monkeypatch.setattr(vault.session, "post", down)
        with pytest.raises(requests.ConnectTimeout):
            vault.add_keys("SVC", {KID: GOOD, SIBLING: BAD})
    assert len(calls) == 2  # no batch step-down, no single-key fallback

    calls.clear()
    vaults = Vaults("SVC")
    vaults.vaults = [api]
    assert vaults.get_key(KID) == (None, None)
    assert vaults.get_key(SIBLING) == (None, None)
    assert vaults.add_keys({KID: GOOD}) == 0
    assert calls == ["https://api.example/svc/" + KID.hex]


def test_a_vault_still_rate_limited_is_not_used_again(monkeypatch: pytest.MonkeyPatch) -> None:
    calls: list[str] = []
    api = API("api", "https://api.example/", "tok")
    monkeypatch.setattr(api.session, "get", lambda url, **kw: calls.append(url) or SimpleNamespace(status_code=429))
    vaults = Vaults("SVC")
    vaults.vaults = [api]

    assert vaults.get_key(KID) == (None, None)
    assert vaults.get_key(SIBLING) == (None, None)
    assert calls == ["https://api.example/svc/" + KID.hex]