mirror of
https://github.com/unshackle-dl/unshackle.git
synced 2026-11-12 12:05:07 +00:00
feat(vaults): flag and skip wrong content keys from poisoned vaults
Verify a vault-sourced key with an FFmpeg decode after decryption. On failure flag the KID:KEY locally with its source, restore the ciphertext, and retry through the remaining vaults and then the CDM.
This commit is contained in:
@@ -248,6 +248,11 @@ Each command reports how many files it removed and how much space it freed, then
|
||||
|
||||
**Inspect:** enable `debug_keys` and operate unshackle with `-d`. Then examine the JSON debug log's `drm_*` and `vault_*` operations to see which content keys unshackle fetched, and from where. Also make sure that shaka-packager (or `mp4decrypt`) shows a sane version in the `download_init` binary-versions entry. See [DRM & CDM](drm-and-cdm.md).
|
||||
|
||||
!!! example "A vault served a wrong content key"
|
||||
**Symptom:** a warning names a content key and the vault it came from, either `was bad, trying other vaults` (the decode check failed) or `is flagged bad, skipping` (an earlier run flagged it). The download stops with `The content key from the vault did not decrypt the track; run again`, `was bad and no other source has the key`, or `No vault or CDM produced a content key that decrypts the track` when no source is left to try.
|
||||
|
||||
**Inspect:** while it decrypts a track whose content key came from a vault, unshackle keeps the ciphertext as a `.enc` sibling of the track file in the temp directory, and removes it afterwards; a leftover one is from a killed run and is safe to delete. Read the `bad_keys` table of your local SQLite vault to see which pairs are flagged and which vault each came from. See [Wrong keys from a poisoned vault](vaults.md#wrong-keys-from-a-poisoned-vault).
|
||||
|
||||
!!! example "A network request fails at the connection or TLS layer"
|
||||
**Symptom:** timeouts, connection resets, or handshake errors, and unshackle's own logs do not give the full cause.
|
||||
|
||||
|
||||
+21
-1
@@ -32,6 +32,22 @@ When unshackle recovers new keys from a license, it pushes them to **all** confi
|
||||
|
||||
unshackle pushes keys from every source, not only a local CDM. A run against a remote server that licenses with its own CDM (`--remote`) writes the keys the server returns, and `unshackle import` writes the keys it reads from an export file.
|
||||
|
||||
### Wrong keys from a poisoned vault
|
||||
|
||||
A public vault can return the right KID with a wrong content key. The decrypters accept a wrong content key and write noise. unshackle therefore trusts a content key from a vault only after FFmpeg decodes the first seconds of the decrypted track with strict error detection. Content keys from a licence skip this check. A file this FFmpeg build cannot judge counts as a pass, never as a wrong content key: a codec it has no decoder for, or a header it cannot read. Without FFmpeg there is no check at all.
|
||||
|
||||
When the decode fails, unshackle:
|
||||
|
||||
1. Flags the KID:KEY in the `bad_keys` table of every local SQLite vault, with the service and the name of the vault it came from. The pair also goes out of the service table, except in a vault marked `no_push`, where only the flag is written.
|
||||
2. Restores the ciphertext from a hard link it made before decryption, then asks for the content key again. Every vault lookup skips a flagged pair, so the request reaches the remaining vaults and then the CDM. A remote answer that arrived after the winning one is reused instead of queried again.
|
||||
3. Removes the flag when the CDM returns the same content key, because then the decode check was wrong, not the content key. A content key from another vault is checked in its turn, and flagged as well if it also fails.
|
||||
|
||||
The download stops with an error instead of writing noise when no ciphertext is left to retry with, or when every vault and the CDM are used up. Without a local SQLite vault nothing remembers a flag, so the same vault serves the same pair again and the download stops there.
|
||||
|
||||
Per-segment decryption gives up the retry. With `decrypt_segments` on, and always with MonaLisa, the plaintext replaces each segment as it arrives, so no ciphertext is left. A failed decode then flags the pair and stops the download. The next run skips the flagged pair.
|
||||
|
||||
unshackle copies a content key from one vault to the others only after this check passes, so a wrong content key does not spread. Without FFmpeg, a content key from a vault is not copied at all. To see or clear the flags, open the SQLite file and read or delete rows in `bad_keys`.
|
||||
|
||||
!!! tip "A good two-vault setup"
|
||||
A common arrangement is a fast local **SQLite** vault plus a shared remote vault (**MySQL**, **HTTP**, or **API**). The local vault answers instantly and works offline; the remote vault lets you share keys across machines or with a group.
|
||||
|
||||
@@ -91,6 +107,8 @@ key_vaults:
|
||||
|
||||
Each service gets its own table (named after the service tag), created on demand with `kid` and `key_` columns. Connections use WAL journaling with a 30-second busy timeout, so multiple threads can safely use the same database file.
|
||||
|
||||
One more table, `bad_keys`, is created on demand and holds every KID:KEY that failed the decode check described in [Wrong keys from a poisoned vault](#wrong-keys-from-a-poisoned-vault), with the service, the vault it came from, and the time it was flagged. It is not a service table, and every command that walks the service tables skips it.
|
||||
|
||||
!!! warning "Moving the .db file silently forks your vault"
|
||||
Because the file is created on demand, if you relocate an existing `.db` without updating `path` to match, unshackle will create a fresh empty database at the old location. That leaves you with two divergent vaults and no error to warn you. Always update the config path whenever you move the file.
|
||||
|
||||
@@ -299,6 +317,8 @@ $ unshackle kv copy local shared community
|
||||
|
||||
`--service` and `--local-only` are mutually exclusive.
|
||||
|
||||
A pair the destination vault has flagged in `bad_keys` is left behind, so a copy does not undo a flag. `bad_keys` is never treated as a service, so its rows are never copied into a service table.
|
||||
|
||||
### `kv sync`: mirror vaults both ways
|
||||
|
||||
Make each of two or more vaults hold every content key the others have. This is effectively a two-way `copy` between each pair, so afterwards all listed vaults hold the same set of keys.
|
||||
@@ -307,7 +327,7 @@ Make each of two or more vaults hold every content key the others have. This is
|
||||
$ unshackle kv sync local shared
|
||||
```
|
||||
|
||||
It accepts the same `--service` and `--local-only` options as `copy`, and requires at least two vaults.
|
||||
It accepts the same `--service` and `--local-only` options as `copy`, and requires at least two vaults. It runs `copy` between each pair, so it skips flagged pairs and the `bad_keys` table in the same way.
|
||||
|
||||
### `kv prepare`: pre-create service tables
|
||||
|
||||
|
||||
@@ -0,0 +1,298 @@
|
||||
import time
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
from uuid import UUID
|
||||
|
||||
import pytest
|
||||
|
||||
from unshackle.commands import dl as dl_module
|
||||
from unshackle.commands.dl import dl
|
||||
from unshackle.core.vaults import Vaults
|
||||
from unshackle.vaults.SQLite import SQLite
|
||||
|
||||
KID = UUID("11111111-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: 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
|
||||
|
||||
local = SQLite("local", tmp_path / "local.db")
|
||||
fast = FakeVault("fast", BAD, delay=0.01)
|
||||
slow = FakeVault("slow", GOOD, delay=0.2)
|
||||
vaults = Vaults("SVC")
|
||||
vaults.vaults = [local, fast, slow]
|
||||
|
||||
assert vaults.get_key(KID) == (BAD, fast)
|
||||
vaults.flag_bad_key(KID, BAD)
|
||||
time.sleep(0.4) # the slow answer lands in candidates after the first call returned
|
||||
assert vaults.get_key(KID) == (GOOD, slow)
|
||||
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:
|
||||
"""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: 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)
|
||||
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_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
|
||||
@@ -29,7 +29,9 @@ KEY = "a" * 32
|
||||
def tables(path: Path) -> list[str]:
|
||||
conn = sqlite3.connect(path)
|
||||
try:
|
||||
return sorted(r[0] for r in conn.execute("SELECT name FROM sqlite_master WHERE type='table'"))
|
||||
return sorted(
|
||||
r[0] for r in conn.execute("SELECT name FROM sqlite_master WHERE type='table' AND name != 'bad_keys'")
|
||||
)
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
@@ -122,11 +122,11 @@ def test_batch_resolved_keys_are_cached_before_download():
|
||||
assert vaults.added == {NEW: "k_new"}
|
||||
|
||||
|
||||
def test_vault_hits_replicate_to_the_other_vaults():
|
||||
"""The local path pushes a vault hit to every other vault; the server path must too."""
|
||||
def test_vault_hits_wait_for_the_decrypt_before_replicating():
|
||||
"""A vault hit reaches the other vaults only after decrypt_verified proves it, never from prepare_drm."""
|
||||
vaults = FakeVaults({CACHED: "k_cached"}, vault_used="vault-a")
|
||||
run_prepare(vaults, [CACHED], {}, track_kid=CACHED)
|
||||
assert vaults.replicated == [(CACHED, "k_cached", "vault-a")]
|
||||
assert vaults.replicated == []
|
||||
assert vaults.added == {}
|
||||
|
||||
|
||||
|
||||
+119
-44
@@ -48,7 +48,7 @@ from unshackle.core.console import GradientPulseBarColumn, SyncLive, console, li
|
||||
from unshackle.core.constants import DOWNLOAD_CANCELLED, DOWNLOAD_LICENCE_ONLY, AnyTrack, context_settings
|
||||
from unshackle.core.credential import Credential
|
||||
from unshackle.core.downloaders import default_max_workers, format_speed, parse_speed_limit, set_speed_limit
|
||||
from unshackle.core.drm import DRM_T, ClearKeyCENC, MonaLisa, PlayReady, Widevine
|
||||
from unshackle.core.drm import DRM_T, ClearKeyCENC, MonaLisa, PlayReady, Widevine, verify
|
||||
from unshackle.core.events import events
|
||||
from unshackle.core.providers.anilist import parse_anilist_ref
|
||||
from unshackle.core.providers.tvdb import SEASON_TYPES, parse_int
|
||||
@@ -102,7 +102,8 @@ from unshackle.core.utils.collections import ci_get, merge_dict
|
||||
from unshackle.core.utils.post_scripts import NO_POST_SCRIPTS, build_context, dispatch, season_context
|
||||
from unshackle.core.utils.redact import mask_proxy
|
||||
from unshackle.core.utils.selector import select_multiple
|
||||
from unshackle.core.utils.subprocess import ffprobe
|
||||
from unshackle.core.utils.subprocess import ffmpeg_decodes, ffprobe
|
||||
from unshackle.core.vault import Vault
|
||||
from unshackle.core.vaults import Vaults
|
||||
|
||||
|
||||
@@ -1233,6 +1234,7 @@ class dl:
|
||||
|
||||
if cdm_only:
|
||||
self.vaults = Vaults(self.vault_service)
|
||||
verify.DECRYPT_HOOK = self.decrypt_verified
|
||||
self.log.info("CDM-only mode: Skipping vault loading")
|
||||
log_event(
|
||||
"vault_loading_skipped",
|
||||
@@ -1243,6 +1245,7 @@ class dl:
|
||||
else:
|
||||
with console.status("Loading Key Vaults...", spinner="dots"):
|
||||
self.vaults = Vaults(self.vault_service)
|
||||
verify.DECRYPT_HOOK = self.decrypt_verified
|
||||
total_vaults = len(config.key_vaults)
|
||||
failed_vaults = []
|
||||
|
||||
@@ -2874,6 +2877,20 @@ class dl:
|
||||
download_table = Table.grid()
|
||||
download_table.add_row(selected_tracks)
|
||||
|
||||
def prepare_drm_for(track: AnyTrack) -> Callable:
|
||||
return partial(
|
||||
partial(self.prepare_drm, table=download_table),
|
||||
track=track,
|
||||
title=title,
|
||||
certificate=partial(service.get_widevine_service_certificate, title=title, track=track),
|
||||
licence=partial(self.service_licence, service, title=title, track=track),
|
||||
clearkey_licence=partial(service.get_clearkey_license, title=title, track=track),
|
||||
cdm_only=cdm_only,
|
||||
vaults_only=vaults_only,
|
||||
export=export_path,
|
||||
service_session=service.session,
|
||||
)
|
||||
|
||||
if getattr(self._remote_service, "_server_cdm", False) and self.cdm is None:
|
||||
cdm_type = getattr(self._remote_service, "_server_cdm_type", "widevine")
|
||||
self.cdm = cdm_type_stub(cdm_type)
|
||||
@@ -2915,31 +2932,7 @@ class dl:
|
||||
track.download(
|
||||
session=track.session or service.session,
|
||||
no_proxy_download=no_proxy_download,
|
||||
prepare_drm=partial(
|
||||
partial(self.prepare_drm, table=download_table),
|
||||
track=track,
|
||||
title=title,
|
||||
certificate=partial(
|
||||
service.get_widevine_service_certificate,
|
||||
title=title,
|
||||
track=track,
|
||||
),
|
||||
licence=partial(
|
||||
self.service_licence,
|
||||
service,
|
||||
title=title,
|
||||
track=track,
|
||||
),
|
||||
clearkey_licence=partial(
|
||||
service.get_clearkey_license,
|
||||
title=title,
|
||||
track=track,
|
||||
),
|
||||
cdm_only=cdm_only,
|
||||
vaults_only=vaults_only,
|
||||
export=export_path,
|
||||
service_session=service.session,
|
||||
),
|
||||
prepare_drm=prepare_drm_for(track),
|
||||
cdm=self.cdm,
|
||||
max_workers=workers,
|
||||
adaptive_workers=adaptive_workers,
|
||||
@@ -3119,7 +3112,7 @@ class dl:
|
||||
for track in drm_tracks:
|
||||
drm = track.get_drm_for_cdm(self.cdm)
|
||||
if drm and hasattr(drm, "decrypt"):
|
||||
drm.decrypt(track.path)
|
||||
self.decrypt_verified(drm, track.path, prepare_drm_for(track))
|
||||
if not isinstance(drm, MonaLisa):
|
||||
# MonaLisa decrypts per segment; its decrypt() here is a no-op
|
||||
assert_fragments_decrypted(track.path)
|
||||
@@ -3977,6 +3970,86 @@ class dl:
|
||||
|
||||
export.write_text(json.dumps(doc, indent=4, ensure_ascii=False), encoding="utf8")
|
||||
|
||||
def decrypt_verified(
|
||||
self,
|
||||
drm: DRM_T,
|
||||
path: Path,
|
||||
licence: Optional[Callable] = None,
|
||||
track_kid: Optional[UUID] = None,
|
||||
decrypt: bool = True,
|
||||
) -> None:
|
||||
"""Decrypt the file, and when a vault supplied a content key, prove the output decodes.
|
||||
|
||||
A poisoned vault returns the right KID with a wrong content key, and the decrypters
|
||||
accept it without complaint. A content key from a vault is therefore trusted, and only
|
||||
then copied to the other vaults, once FFmpeg can decode the result. Without FFmpeg
|
||||
neither happens. On failure the pair is flagged in the local vaults with
|
||||
its source, the ciphertext is restored, and the licence runs again. Every vault skips a
|
||||
flagged pair, so each pass burns at most one vault; when the vaults run out the CDM
|
||||
answers, and a CDM key is not checked. A CDM that returns the flagged key clears the
|
||||
flag, because then the decode check was wrong.
|
||||
|
||||
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)
|
||||
if keys is None:
|
||||
if decrypt:
|
||||
drm.decrypt(path)
|
||||
return
|
||||
known_bad = [kid for kid, key in keys.items() if (kid, key) in self.vaults.flagged]
|
||||
if known_bad and licence:
|
||||
for kid in known_bad:
|
||||
keys.pop(kid)
|
||||
self.LICENSE_KEY_CACHE.pop(kid, None)
|
||||
licence(drm, track_kid=track_kid)
|
||||
|
||||
def vault_kids() -> dict[UUID, Vault]:
|
||||
"""The KIDs whose key on the track came from a vault, with that vault."""
|
||||
found = {kid: self.vaults.sources.get(kid) for kid in keys}
|
||||
return {kid: src[1] for kid, src in found.items() if src and src[0] == keys[kid]}
|
||||
|
||||
backup = path.with_name(path.name + ".enc")
|
||||
backup.unlink(missing_ok=True) # a killed run leaves one behind; its bytes may not match
|
||||
if vault_kids() and decrypt and binaries.FFMPEG:
|
||||
try:
|
||||
os.link(path, backup)
|
||||
except OSError as e:
|
||||
self.log.debug(f"Cannot keep the ciphertext for a decrypt retry: {e!r}")
|
||||
try:
|
||||
for _ in range(len(self.vaults) + 1):
|
||||
if decrypt:
|
||||
drm.decrypt(path)
|
||||
kids = vault_kids()
|
||||
if not kids or not binaries.FFMPEG or not path.exists():
|
||||
return
|
||||
if ffmpeg_decodes(path):
|
||||
self.flush_vault_writes(
|
||||
[partial(self.vaults.add_key, kid, keys[kid], excluding=vault) for kid, vault in kids.items()]
|
||||
)
|
||||
return
|
||||
bad = {kid: keys.pop(kid) for kid in kids}
|
||||
for kid, key in bad.items():
|
||||
self.log.warning(f"{key} from {kids[kid].name} was bad, trying other vaults")
|
||||
self.vaults.flag_bad_key(kid, key)
|
||||
self.LICENSE_KEY_CACHE.pop(kid, None)
|
||||
if not backup.exists() or not licence:
|
||||
raise ValueError("The content key from the vault did not decrypt the track; run again")
|
||||
path.unlink()
|
||||
os.link(backup, path)
|
||||
licence(drm, track_kid=track_kid)
|
||||
for kid, key in bad.items():
|
||||
if keys.get(kid) != key:
|
||||
continue
|
||||
source = self.vaults.sources.get(kid)
|
||||
if source is None:
|
||||
self.vaults.unflag_bad_key(kid, key)
|
||||
elif source[1] is kids[kid]:
|
||||
raise ValueError(f"{key} from {kids[kid].name} was bad and no other source has the key")
|
||||
raise ValueError("No vault or CDM produced a content key that decrypts the track")
|
||||
finally:
|
||||
backup.unlink(missing_ok=True)
|
||||
|
||||
def flush_vault_writes(self, pending: list[Callable[[], Any]]) -> None:
|
||||
"""Hand the queued vault writes to the background writer so the track does not wait on them."""
|
||||
for write in pending:
|
||||
@@ -4090,10 +4163,6 @@ class dl:
|
||||
content_key = self.LICENSE_KEY_CACHE.get(kid)
|
||||
if not content_key and not cdm_only:
|
||||
content_key, vault_used = self.vaults.get_key(kid)
|
||||
if content_key:
|
||||
pending_vault_writes.append(
|
||||
partial(self.vaults.add_key, kid, content_key, excluding=vault_used)
|
||||
)
|
||||
if content_key:
|
||||
drm.content_keys[kid] = content_key
|
||||
self.LICENSE_KEY_CACHE[kid] = content_key
|
||||
@@ -4278,9 +4347,6 @@ class dl:
|
||||
label = f"[text2]{kid.hex}:{content_key}{is_track_kid} from {vault_used}"
|
||||
if not any(f"{kid.hex}:{content_key}" in x.label for x in cek_tree.children):
|
||||
cek_tree.add(label)
|
||||
pending_vault_writes.append(
|
||||
partial(self.vaults.add_key, kid, content_key, excluding=vault_used)
|
||||
)
|
||||
self.LICENSE_KEY_CACHE[kid] = content_key
|
||||
|
||||
if self.debug_logger:
|
||||
@@ -4384,7 +4450,12 @@ class dl:
|
||||
|
||||
self.LICENSE_KEY_CACHE.update(drm.content_keys)
|
||||
|
||||
pending_vault_writes.append(partial(self.cache_keys_to_vaults, dict(drm.content_keys)))
|
||||
pending_vault_writes.append(
|
||||
partial(
|
||||
self.cache_keys_to_vaults,
|
||||
{k: v for k, v in drm.content_keys.items() if k not in from_vaults},
|
||||
)
|
||||
)
|
||||
|
||||
if track_kid and track_kid not in drm.content_keys:
|
||||
msg = f"No Content Key for KID {track_kid.hex} was returned in the License"
|
||||
@@ -4475,9 +4546,6 @@ class dl:
|
||||
label = f"[text2]{kid.hex}:{content_key}{is_track_kid} from {vault_used}"
|
||||
if not any(f"{kid.hex}:{content_key}" in x.label for x in cek_tree.children):
|
||||
cek_tree.add(label)
|
||||
pending_vault_writes.append(
|
||||
partial(self.vaults.add_key, kid, content_key, excluding=vault_used)
|
||||
)
|
||||
self.LICENSE_KEY_CACHE[kid] = content_key
|
||||
|
||||
if self.debug_logger:
|
||||
@@ -4556,7 +4624,12 @@ class dl:
|
||||
|
||||
self.LICENSE_KEY_CACHE.update(drm.content_keys)
|
||||
|
||||
pending_vault_writes.append(partial(self.cache_keys_to_vaults, dict(drm.content_keys)))
|
||||
pending_vault_writes.append(
|
||||
partial(
|
||||
self.cache_keys_to_vaults,
|
||||
{k: v for k, v in drm.content_keys.items() if k not in from_vaults},
|
||||
)
|
||||
)
|
||||
|
||||
if track_kid and track_kid not in drm.content_keys:
|
||||
msg = f"No Content Key for KID {track_kid.hex} was returned in the License"
|
||||
@@ -4626,9 +4699,6 @@ class dl:
|
||||
label = f"[text2]{kid.hex}:{content_key}{is_track_kid} from {vault_used}"
|
||||
if not any(f"{kid.hex}:{content_key}" in x.label for x in cek_tree.children):
|
||||
cek_tree.add(label)
|
||||
pending_vault_writes.append(
|
||||
partial(self.vaults.add_key, kid, content_key, excluding=vault_used)
|
||||
)
|
||||
self.LICENSE_KEY_CACHE[kid] = content_key
|
||||
elif vaults_only:
|
||||
msg = f"No Vault has a Key for {kid.hex} and --vaults-only was used"
|
||||
@@ -4695,7 +4765,12 @@ class dl:
|
||||
|
||||
self.LICENSE_KEY_CACHE.update(drm.content_keys)
|
||||
|
||||
pending_vault_writes.append(partial(self.cache_keys_to_vaults, dict(drm.content_keys)))
|
||||
pending_vault_writes.append(
|
||||
partial(
|
||||
self.cache_keys_to_vaults,
|
||||
{k: v for k, v in drm.content_keys.items() if k not in from_vaults},
|
||||
)
|
||||
)
|
||||
|
||||
if track_kid and track_kid not in drm.content_keys:
|
||||
msg = f"No Content Key for KID {track_kid.hex} was returned in the License"
|
||||
|
||||
@@ -101,12 +101,15 @@ def add_keys_with_progress(vault: Vault, service: str, kid_keys: dict[str, str],
|
||||
|
||||
def copy_service_data(to_vault: Vault, from_vault: Vault, service: str, log: logging.Logger) -> int:
|
||||
"""Copy data for a single service between vaults."""
|
||||
if service.lower() == "bad_keys":
|
||||
return 0
|
||||
try:
|
||||
content_keys = process_service_keys(from_vault, service, log)
|
||||
except Exception as e:
|
||||
log.warning(f"{service}: Could not read from {from_vault} ({e}), skipped")
|
||||
return 0
|
||||
|
||||
content_keys = {kid: key for kid, key in content_keys.items() if not to_vault.is_bad_key(kid, key)}
|
||||
total_count = len(content_keys)
|
||||
|
||||
if total_count == 0:
|
||||
|
||||
@@ -0,0 +1,23 @@
|
||||
from pathlib import Path
|
||||
from typing import Any, Callable, Optional
|
||||
from uuid import UUID
|
||||
|
||||
DECRYPT_HOOK: Optional[Callable[..., None]] = None
|
||||
|
||||
|
||||
def decrypt_track(
|
||||
drm: Any,
|
||||
path: Path,
|
||||
licence: Optional[Callable] = None,
|
||||
track_kid: Optional[UUID] = None,
|
||||
decrypt: bool = True,
|
||||
) -> None:
|
||||
"""Decrypt ``path`` with ``drm`` through the installed hook, or directly when there is none.
|
||||
|
||||
``decrypt=False`` means the downloader already decrypted the segments in place; the hook
|
||||
then only checks the result, because there is no ciphertext left to retry with.
|
||||
"""
|
||||
if DECRYPT_HOOK:
|
||||
DECRYPT_HOOK(drm, path, licence, track_kid, decrypt)
|
||||
elif decrypt:
|
||||
drm.decrypt(path)
|
||||
@@ -27,6 +27,7 @@ from unshackle.core.config import config
|
||||
from unshackle.core.constants import DOWNLOAD_CANCELLED, DOWNLOAD_LICENCE_ONLY, AnyTrack
|
||||
from unshackle.core.drm import DRM_T, ClearKeyCENC, PlayReady, Widevine
|
||||
from unshackle.core.drm.segment_decrypt import SegmentDecrypter, can_use
|
||||
from unshackle.core.drm.verify import decrypt_track
|
||||
from unshackle.core.events import events
|
||||
from unshackle.core.session import RnetSession
|
||||
from unshackle.core.tracks import Audio, DownloadContext, Subtitle, Tracks, Video, resume
|
||||
@@ -643,8 +644,7 @@ class DASH:
|
||||
|
||||
if drm:
|
||||
progress(downloaded="Decrypting", completed=0, total=None)
|
||||
if not decrypter:
|
||||
drm.decrypt(save_path)
|
||||
decrypt_track(drm, save_path, license_widevine, decrypt=not decrypter)
|
||||
assert_fragments_decrypted(save_path)
|
||||
track.drm = None
|
||||
events.emit(events.Types.TRACK_DECRYPTED, track=track, drm=drm, segment=None)
|
||||
|
||||
@@ -32,6 +32,7 @@ from unshackle.core.config import config
|
||||
from unshackle.core.constants import DOWNLOAD_CANCELLED, DOWNLOAD_LICENCE_ONLY, AnyTrack
|
||||
from unshackle.core.drm import DRM_T, ClearKey, MonaLisa, PlayReady, Widevine
|
||||
from unshackle.core.drm.segment_decrypt import SegmentDecrypter, can_use
|
||||
from unshackle.core.drm.verify import decrypt_track
|
||||
from unshackle.core.events import events
|
||||
from unshackle.core.session import RnetResponse, RnetSession
|
||||
from unshackle.core.tracks import Audio, DownloadContext, Subtitle, Tracks, Video, resume
|
||||
@@ -965,8 +966,7 @@ class HLS:
|
||||
decrypted_init = segment_decrypter.finish()
|
||||
map_data = (map_data[0], decrypted_init)
|
||||
merge(to=merged_path, via=files, delete=True, include_map_data=True)
|
||||
if not segment_decrypter:
|
||||
drm.decrypt(merged_path)
|
||||
decrypt_track(drm, merged_path, license_widevine, decrypt=not segment_decrypter)
|
||||
assert_fragments_decrypted(merged_path)
|
||||
merged_path.rename(decrypted_path)
|
||||
else:
|
||||
|
||||
@@ -20,6 +20,7 @@ from requests import Session
|
||||
from unshackle.core.config import config
|
||||
from unshackle.core.constants import DOWNLOAD_CANCELLED, DOWNLOAD_LICENCE_ONLY, AnyTrack
|
||||
from unshackle.core.drm import DRM_T, PlayReady, Widevine
|
||||
from unshackle.core.drm.verify import decrypt_track
|
||||
from unshackle.core.events import events
|
||||
from unshackle.core.manifests.ism_init import (
|
||||
build_init_segment,
|
||||
@@ -574,7 +575,7 @@ class ISM:
|
||||
|
||||
if session_drm:
|
||||
progress(downloaded="Decrypting", completed=0, total=None)
|
||||
session_drm.decrypt(save_path)
|
||||
decrypt_track(session_drm, save_path, license_widevine)
|
||||
assert_fragments_decrypted(save_path)
|
||||
track.drm = None
|
||||
events.emit(events.Types.TRACK_DECRYPTED, track=track, drm=session_drm, segment=None)
|
||||
|
||||
@@ -24,6 +24,7 @@ from unshackle.core.config import config
|
||||
from unshackle.core.constants import DOWNLOAD_CANCELLED, DOWNLOAD_LICENCE_ONLY, DownloadCancelled
|
||||
from unshackle.core.downloaders import requests
|
||||
from unshackle.core.drm import DRM_T, ClearKeyCENC, PlayReady, Widevine
|
||||
from unshackle.core.drm.verify import decrypt_track
|
||||
from unshackle.core.events import events
|
||||
from unshackle.core.session import RnetSession
|
||||
from unshackle.core.tracks import resume
|
||||
@@ -717,7 +718,7 @@ class Track:
|
||||
|
||||
if drm:
|
||||
progress(downloaded="Decrypting", completed=0, total=None)
|
||||
drm.decrypt(save_path)
|
||||
decrypt_track(drm, save_path, prepare_drm, track_kid)
|
||||
assert_fragments_decrypted(save_path)
|
||||
self.drm = None
|
||||
events.emit(events.Types.TRACK_DECRYPTED, track=self, drm=drm, segment=None)
|
||||
|
||||
@@ -40,6 +40,42 @@ def log_tool_run(
|
||||
)
|
||||
|
||||
|
||||
def ffmpeg_decodes(path: Path, seconds: int = 3) -> bool:
|
||||
"""Return True when FFmpeg decodes the first seconds of the file without an error.
|
||||
|
||||
A wrong content key leaves the container intact and the samples as noise, which the
|
||||
decoders reject. ``-xerror`` stops at the first error so the check stays short.
|
||||
A file this FFmpeg build cannot judge counts as a pass, never as a wrong key.
|
||||
"""
|
||||
if not binaries.FFMPEG:
|
||||
raise EnvironmentError('FFmpeg executable "ffmpeg" not found but is required.')
|
||||
args: list[str] = [
|
||||
str(binaries.FFMPEG),
|
||||
"-nostdin",
|
||||
"-v",
|
||||
"error",
|
||||
"-err_detect",
|
||||
"explode",
|
||||
"-xerror",
|
||||
"-i",
|
||||
str(path),
|
||||
"-t",
|
||||
str(seconds),
|
||||
"-f",
|
||||
"null",
|
||||
"-",
|
||||
]
|
||||
ff = subprocess.run(args, capture_output=True, text=True, encoding="utf-8", errors="replace")
|
||||
if ff.returncode == 0:
|
||||
return True
|
||||
err = ff.stderr.lower()
|
||||
if any(
|
||||
m in err for m in ("not found for", "no decoder found", "sub-sample encryption info", "error reading header")
|
||||
):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def ffprobe(uri: Union[bytes, Path]) -> dict:
|
||||
"""Use FFprobe on the provided data to get its track information (``-show_streams``)."""
|
||||
if not binaries.FFProbe:
|
||||
|
||||
@@ -53,5 +53,15 @@ class Vault(metaclass=ABCMeta):
|
||||
get_key/get_keys/add_key/add_keys reaches the namespace it came from.
|
||||
"""
|
||||
|
||||
def is_bad_key(self, kid: Union[UUID, str], key: str) -> bool:
|
||||
"""Return True if this Vault has flagged the KID:KEY as wrong."""
|
||||
return False
|
||||
|
||||
def flag_bad_key(self, service: str, kid: Union[UUID, str], key: str, source: str) -> None:
|
||||
"""Remember a KID:KEY that failed decryption and where it came from."""
|
||||
|
||||
def unflag_bad_key(self, kid: Union[UUID, str], key: str) -> None:
|
||||
"""Forget a flagged KID:KEY, for example when a CDM licence confirmed it."""
|
||||
|
||||
|
||||
__all__ = ("Vault",)
|
||||
|
||||
@@ -1,7 +1,8 @@
|
||||
import inspect
|
||||
import logging
|
||||
import time
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
from concurrent.futures import Future, ThreadPoolExecutor, as_completed
|
||||
from functools import partial
|
||||
from typing import Any, Iterator, Optional, Union
|
||||
from uuid import UUID
|
||||
|
||||
@@ -24,6 +25,9 @@ class Vaults:
|
||||
def __init__(self, service: Optional[str] = None):
|
||||
self.service = service or ""
|
||||
self.vaults = []
|
||||
self.sources: dict[Union[UUID, str], tuple[str, Vault]] = {}
|
||||
self.flagged: set[tuple[Union[UUID, str], str]] = set()
|
||||
self.candidates: dict[Union[UUID, str], list[tuple[str, Vault]]] = {}
|
||||
|
||||
def __iter__(self) -> Iterator[Vault]:
|
||||
return iter(self.vaults)
|
||||
@@ -60,30 +64,56 @@ class Vaults:
|
||||
"""Get Key from the first Vault it can by KID (Key ID) and Service.
|
||||
|
||||
Local vaults go first, one at a time: a hit there costs nothing and skips the network.
|
||||
The remaining vaults run at the same time, and the first one with the content key wins.
|
||||
A remote answer kept from an earlier call comes next, so a retry after a flagged pair
|
||||
reuses it instead of querying again. The remaining vaults then run at the same time,
|
||||
and the first one with the content key wins.
|
||||
"""
|
||||
local = [v for v in self.vaults if v.local]
|
||||
remote = [v for v in self.vaults if not v.local]
|
||||
for vault in local:
|
||||
key = self.query_vault(vault, kid)
|
||||
if key:
|
||||
self.sources[kid] = (key, vault)
|
||||
return key, vault
|
||||
for key, vault in self.candidates.get(kid, []):
|
||||
if not self.is_flagged(kid, key):
|
||||
self.sources[kid] = (key, vault)
|
||||
return key, vault
|
||||
if not remote:
|
||||
return None, None
|
||||
pool = ThreadPoolExecutor(len(remote))
|
||||
try:
|
||||
futures = {pool.submit(self.query_vault, vault, kid): vault for vault in remote}
|
||||
for future, vault in futures.items():
|
||||
future.add_done_callback(partial(self.remember_candidate, kid, vault))
|
||||
for future in as_completed(futures):
|
||||
key = future.result()
|
||||
if key:
|
||||
self.sources[kid] = (key, futures[future])
|
||||
return key, futures[future]
|
||||
return None, None
|
||||
finally:
|
||||
# a slow vault must not hold up a content key another vault already returned
|
||||
# a slow vault must not hold up a content key another vault already returned; the
|
||||
# ones already running still finish and land in candidates through the callback
|
||||
pool.shutdown(wait=False, cancel_futures=True)
|
||||
|
||||
def remember_candidate(self, kid: Union[UUID, str], vault: Vault, future: Future) -> None:
|
||||
"""Keep a remote answer for a later retry. Cancelled or failed queries leave nothing."""
|
||||
if future.cancelled() or future.exception():
|
||||
return
|
||||
key = future.result()
|
||||
if key and (key, vault) not in self.candidates.setdefault(kid, []):
|
||||
self.candidates[kid].append((key, vault))
|
||||
|
||||
def is_flagged(self, kid: Union[UUID, str], key: str) -> bool:
|
||||
"""True when a local vault has flagged the KID:KEY as wrong."""
|
||||
return any(v.is_bad_key(kid, key) for v in self.vaults if v.local)
|
||||
|
||||
def query_vault(self, vault: Vault, kid: Union[UUID, str]) -> Optional[str]:
|
||||
"""Ask one vault for a content key. A failure logs a warning and returns None."""
|
||||
"""Ask one vault for a content key. A failure logs a warning and returns None.
|
||||
|
||||
A flagged pair counts as no content key, so no vault can serve one again.
|
||||
"""
|
||||
dl = get_debug_logger()
|
||||
start = time.monotonic()
|
||||
try:
|
||||
@@ -103,6 +133,9 @@ class Vaults:
|
||||
)
|
||||
return None
|
||||
found = bool(key and key.count("0") != len(key))
|
||||
if found and key and self.is_flagged(kid, key):
|
||||
log.warning(f"{key} from {vault.name} is flagged bad, skipping")
|
||||
found = False
|
||||
if dl:
|
||||
dl.log_vault_query(
|
||||
vault.name,
|
||||
@@ -114,6 +147,21 @@ class Vaults:
|
||||
)
|
||||
return key if found else None
|
||||
|
||||
def flag_bad_key(self, kid: Union[UUID, str], key: str) -> None:
|
||||
"""Flag a KID:KEY that failed decryption in every local Vault, naming the Vault it came from."""
|
||||
self.flagged.add((kid, key))
|
||||
source = self.sources.pop(kid, None)
|
||||
for vault in self.vaults:
|
||||
if vault.local:
|
||||
vault.flag_bad_key(self.service, kid, key, source[1].name if source else "unknown")
|
||||
|
||||
def unflag_bad_key(self, kid: Union[UUID, str], key: str) -> None:
|
||||
"""Remove a flag from every local Vault."""
|
||||
self.flagged.discard((kid, key))
|
||||
for vault in self.vaults:
|
||||
if vault.local:
|
||||
vault.unflag_bad_key(kid, key)
|
||||
|
||||
def add_key(self, kid: Union[UUID, str], key: str, excluding: Optional[Vault] = None) -> int:
|
||||
"""Add a KID:KEY to all Vaults, optionally with an exclusion.
|
||||
|
||||
|
||||
@@ -27,11 +27,16 @@ class SQLite(Vault):
|
||||
if not table:
|
||||
return None
|
||||
|
||||
self.create_bad_keys_table()
|
||||
conn = self.conn_factory.get()
|
||||
cursor = conn.cursor()
|
||||
|
||||
try:
|
||||
cursor.execute(f"SELECT `id`, `key_` FROM `{table}` WHERE `kid`=? AND `key_`!=?", (kid, "0" * 32))
|
||||
cursor.execute(
|
||||
f"SELECT `id`, `key_` FROM `{table}` WHERE `kid`=? AND `key_`!=? "
|
||||
"AND `key_` NOT IN (SELECT `key_` FROM `bad_keys` WHERE `kid`=?)",
|
||||
(kid, "0" * 32, kid),
|
||||
)
|
||||
cek = cursor.fetchone()
|
||||
return cek[1] if cek else None
|
||||
finally:
|
||||
@@ -150,11 +155,77 @@ class SQLite(Vault):
|
||||
try:
|
||||
cursor.execute("SELECT name FROM sqlite_master WHERE type='table';")
|
||||
for (name,) in cursor.fetchall():
|
||||
if name != "sqlite_sequence":
|
||||
if name not in ("sqlite_sequence", "bad_keys"):
|
||||
yield name
|
||||
finally:
|
||||
cursor.close()
|
||||
|
||||
def is_bad_key(self, kid: Union[UUID, str], key: str) -> bool:
|
||||
if isinstance(kid, UUID):
|
||||
kid = kid.hex
|
||||
self.create_bad_keys_table()
|
||||
cursor = self.conn_factory.get().cursor()
|
||||
try:
|
||||
cursor.execute("SELECT 1 FROM `bad_keys` WHERE `kid`=? AND `key_`=?", (kid, key))
|
||||
return cursor.fetchone() is not None
|
||||
finally:
|
||||
cursor.close()
|
||||
|
||||
def flag_bad_key(self, service: str, kid: Union[UUID, str], key: str, source: str) -> None:
|
||||
if isinstance(kid, UUID):
|
||||
kid = kid.hex
|
||||
self.create_bad_keys_table()
|
||||
conn = self.conn_factory.get()
|
||||
cursor = conn.cursor()
|
||||
try:
|
||||
cursor.execute(
|
||||
"INSERT OR IGNORE INTO `bad_keys` (service, kid, key_, source) VALUES (?, ?, ?, ?)",
|
||||
(service, kid, key, source),
|
||||
)
|
||||
table = None if self.no_push else self.resolve_table(service)
|
||||
if table:
|
||||
cursor.execute(f"DELETE FROM `{table}` WHERE `kid`=? AND `key_`=?", (kid, key))
|
||||
finally:
|
||||
conn.commit()
|
||||
cursor.close()
|
||||
|
||||
def unflag_bad_key(self, kid: Union[UUID, str], key: str) -> None:
|
||||
if isinstance(kid, UUID):
|
||||
kid = kid.hex
|
||||
self.create_bad_keys_table()
|
||||
conn = self.conn_factory.get()
|
||||
cursor = conn.cursor()
|
||||
try:
|
||||
cursor.execute("DELETE FROM `bad_keys` WHERE `kid`=? AND `key_`=?", (kid, key))
|
||||
finally:
|
||||
conn.commit()
|
||||
cursor.close()
|
||||
|
||||
def create_bad_keys_table(self) -> None:
|
||||
"""Make the bad_keys table if it does not exist yet. Runs the DDL once per instance."""
|
||||
if getattr(self, "_bad_keys_ready", False):
|
||||
return
|
||||
conn = self.conn_factory.get()
|
||||
cursor = conn.cursor()
|
||||
try:
|
||||
cursor.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS bad_keys (
|
||||
"service" TEXT NOT NULL,
|
||||
"kid" TEXT NOT NULL COLLATE NOCASE,
|
||||
"key_" TEXT NOT NULL COLLATE NOCASE,
|
||||
"source" TEXT NOT NULL,
|
||||
"flagged_at" TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
UNIQUE("kid", "key_")
|
||||
);
|
||||
"""
|
||||
)
|
||||
finally:
|
||||
conn.commit()
|
||||
cursor.close()
|
||||
# set only after the commit: a parallel vault thread must not query a table still being made
|
||||
self._bad_keys_ready = True
|
||||
|
||||
def resolve_table(self, name: str) -> Optional[str]:
|
||||
"""Get the actual Table name matching `name` case-insensitively, if it exists."""
|
||||
conn = self.conn_factory.get()
|
||||
|
||||
Reference in New Issue
Block a user