mirror of
https://github.com/unshackle-dl/unshackle.git
synced 2026-11-12 20:15:10 +00:00
fix(vaults): stop querying sibling KIDs and skip unavailable vaults
The server CDM stub names only the manifest's KIDs; sibling keys stay content keys. API and HTTP vaults retry a 429, and a vault that is down or still rate limited costs one timeout per run.
This commit is contained in:
@@ -56,3 +56,9 @@ key_vaults:
|
|||||||
Timeout in seconds for vault operations. Injected automatically into any backend whose
|
Timeout in seconds for vault operations. Injected automatically into any backend whose
|
||||||
constructor accepts a `timeout` parameter (a per-vault `timeout` still wins). The `MySQL` backend
|
constructor accepts a `timeout` parameter (a per-vault `timeout` still wins). The `MySQL` backend
|
||||||
uses it as the default `connect_timeout`.
|
uses it as the default `connect_timeout`.
|
||||||
|
|
||||||
|
The `API` and `HTTP` backends send a request again up to 3 times when the vault answers
|
||||||
|
`429 Too Many Requests`. They wait the `Retry-After` time, but never longer than this timeout.
|
||||||
|
A connection failure or a timeout is not tried again. The run does not use a vault again after a
|
||||||
|
connection failure, a timeout, or a `429` that continues after the retries, so a vault that is
|
||||||
|
down costs one timeout per run.
|
||||||
|
|||||||
@@ -10,11 +10,14 @@ import pytest
|
|||||||
|
|
||||||
from unshackle.commands import dl as dl_module
|
from unshackle.commands import dl as dl_module
|
||||||
from unshackle.commands.dl import dl
|
from unshackle.commands.dl import dl
|
||||||
|
from unshackle.core.vault import RateLimited, vault_session
|
||||||
from unshackle.core.vaults import Vaults
|
from unshackle.core.vaults import Vaults
|
||||||
from unshackle.vaults.API import API
|
from unshackle.vaults.API import API, Exceptions
|
||||||
|
from unshackle.vaults.HTTP import HTTP
|
||||||
from unshackle.vaults.SQLite import SQLite
|
from unshackle.vaults.SQLite import SQLite
|
||||||
|
|
||||||
KID = UUID("11111111-2222-3333-4444-555555555555")
|
KID = UUID("11111111-2222-3333-4444-555555555555")
|
||||||
|
SIBLING = UUID("99999999-2222-3333-4444-555555555555")
|
||||||
BAD = "aa" * 16
|
BAD = "aa" * 16
|
||||||
GOOD = "bb" * 16
|
GOOD = "bb" * 16
|
||||||
|
|
||||||
@@ -630,3 +633,100 @@ def test_api_vault_reports_the_pair_upstream(monkeypatch: pytest.MonkeyPatch, ca
|
|||||||
vaults.flag_bad_key(KID, BAD)
|
vaults.flag_bad_key(KID, BAD)
|
||||||
vaults.unflag_bad_key(KID, BAD)
|
vaults.unflag_bad_key(KID, BAD)
|
||||||
assert not caplog.records
|
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))
|
||||||
|
cmd._remote_service = SimpleNamespace(server_vault_keys={})
|
||||||
|
video, audio = (SimpleNamespace(kids=[KID], content_keys={KID: GOOD, SIBLING: BAD}) for _ in range(2))
|
||||||
|
|
||||||
|
cmd.prefer_vault_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]
|
||||||
|
cmd._remote_service = SimpleNamespace(server_vault_keys={})
|
||||||
|
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(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]
|
||||||
|
|||||||
@@ -170,3 +170,51 @@ def test_resolve_server_keys_does_not_warn_for_clear_tracks(monkeypatch):
|
|||||||
svc.resolve_server_keys(SimpleNamespace(tracks=Tracks([video_track, audio_track])))
|
svc.resolve_server_keys(SimpleNamespace(tracks=Tracks([video_track, audio_track])))
|
||||||
|
|
||||||
assert warnings == ["Server CDM returned no content keys for track v-1"]
|
assert warnings == ["Server CDM returned no content keys for track v-1"]
|
||||||
|
|
||||||
|
|
||||||
|
def resolve(manifest_drm, drm_type="widevine"):
|
||||||
|
"""Run resolve_server_keys for one track the server answers with its own key and a sibling."""
|
||||||
|
svc = remote_service()
|
||||||
|
svc.log = SimpleNamespace(warning=lambda *a, **k: None, debug=lambda *a, **k: None)
|
||||||
|
svc._server_cdm = True
|
||||||
|
svc.drain_server_logs = lambda: None
|
||||||
|
keys = {CDM_KID.hex: "bb" * 16, KID.hex: "aa" * 16}
|
||||||
|
svc.client = SimpleNamespace(post=lambda ep, data: {"keys": {"v-1": keys}, "drm_type": drm_type})
|
||||||
|
track = SimpleNamespace(id="v-1", drm=manifest_drm)
|
||||||
|
|
||||||
|
class Tracks(list):
|
||||||
|
videos = [track]
|
||||||
|
audio: list = []
|
||||||
|
|
||||||
|
svc.resolve_server_keys(SimpleNamespace(tracks=Tracks([track])))
|
||||||
|
assert set(track.drm[0].content_keys) == {KID, CDM_KID}
|
||||||
|
return track.drm[0]
|
||||||
|
|
||||||
|
|
||||||
|
def test_resolve_server_keys_names_only_the_manifest_kids():
|
||||||
|
"""Sibling keys in the licence stay content keys but are not KIDs the track names."""
|
||||||
|
from pywidevine.pssh import PSSH
|
||||||
|
|
||||||
|
from unshackle.core.drm import Widevine
|
||||||
|
|
||||||
|
widevine = Widevine(pssh=PSSH.new(system_id=PSSH.SystemId.Widevine, key_ids=[KID]))
|
||||||
|
assert resolve([widevine]).kids == [KID]
|
||||||
|
|
||||||
|
playready = RemoteService.create_drm_stub("playready", [KID.hex])
|
||||||
|
assert resolve([playready], "playready").kids == [KID]
|
||||||
|
|
||||||
|
|
||||||
|
def test_resolve_server_keys_names_every_kid_without_manifest_drm():
|
||||||
|
assert set(resolve(None).kids) == {KID, CDM_KID}
|
||||||
|
|
||||||
|
|
||||||
|
def test_reconstruct_drm_keeps_the_default_kid_of_a_pssh_without_kids():
|
||||||
|
from pywidevine.pssh import PSSH
|
||||||
|
|
||||||
|
from unshackle.core.drm import real_kids
|
||||||
|
from unshackle.core.remote_service import reconstruct_drm
|
||||||
|
|
||||||
|
content_id_only = PSSH.new(system_id=PSSH.SystemId.Widevine, init_data=b"\x22\x03abc")
|
||||||
|
(drm,) = reconstruct_drm([{"type": "widevine", "pssh": content_id_only.dumps(), "kids": [str(KID)]}])
|
||||||
|
|
||||||
|
assert real_kids(drm) == [KID]
|
||||||
|
|||||||
@@ -33,6 +33,9 @@ class FakeVaults:
|
|||||||
def __len__(self):
|
def __len__(self):
|
||||||
return 1
|
return 1
|
||||||
|
|
||||||
|
def held_locally(self, kid, key):
|
||||||
|
return False
|
||||||
|
|
||||||
def get_key(self, kid):
|
def get_key(self, kid):
|
||||||
return self.held.get(kid), self.vault_used
|
return self.held.get(kid), self.vault_used
|
||||||
|
|
||||||
|
|||||||
@@ -55,7 +55,7 @@ from unshackle.core.constants import (
|
|||||||
)
|
)
|
||||||
from unshackle.core.credential import Credential
|
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.downloaders import default_max_workers, format_speed, parse_speed_limit, set_speed_limit
|
||||||
from unshackle.core.drm import DRM_T, ClearKeyCENC, MonaLisa, PlayReady, Widevine, cbcs, own_kids, verify
|
from unshackle.core.drm import DRM_T, ClearKeyCENC, MonaLisa, PlayReady, Widevine, cbcs, own_kids, real_kids, verify
|
||||||
from unshackle.core.events import events
|
from unshackle.core.events import events
|
||||||
from unshackle.core.export_name import ExportNamer, move_export, plural, season_episodes, title_label
|
from unshackle.core.export_name import ExportNamer, move_export, plural, season_episodes, title_label
|
||||||
from unshackle.core.providers.anilist import parse_anilist_ref
|
from unshackle.core.providers.anilist import parse_anilist_ref
|
||||||
@@ -4555,13 +4555,21 @@ class dl:
|
|||||||
poisoned row; the server key waits as the next candidate. A server CDM key is trusted
|
poisoned row; the server key waits as the next candidate. A server CDM key is trusted
|
||||||
like a local CDM key: it goes to the run cache and the vaults, and a retry after a
|
like a local CDM key: it goes to the run cache and the vaults, and a retry after a
|
||||||
flagged vault key takes it before any vault. A server vault key is unproven, so it
|
flagged vault key takes it before any vault. A server vault key is unproven, so it
|
||||||
waits in ``Vaults.candidates`` and gets the same check when its turn comes.
|
waits in ``Vaults.candidates`` and gets the same check when its turn comes. Only the KIDs
|
||||||
|
the DRM names are looked up, each once, so sibling keys add no remote vault requests.
|
||||||
"""
|
"""
|
||||||
server_vault_keys = getattr(self._remote_service, "server_vault_keys", {})
|
server_vault_keys = getattr(self._remote_service, "server_vault_keys", {})
|
||||||
|
looked_up: dict[UUID, tuple[Optional[str], Any]] = {}
|
||||||
for track in title.tracks:
|
for track in title.tracks:
|
||||||
for drm in getattr(track, "drm", None) or []:
|
for drm in getattr(track, "drm", None) or []:
|
||||||
for kid, server_key in list(getattr(drm, "content_keys", {}).items()):
|
content_keys = getattr(drm, "content_keys", {})
|
||||||
vault_key, vault = self.vaults.get_key(kid)
|
named = set(real_kids(drm) or content_keys)
|
||||||
|
for kid, server_key in list(content_keys.items()):
|
||||||
|
if kid not in named:
|
||||||
|
continue
|
||||||
|
if kid not in looked_up:
|
||||||
|
looked_up[kid] = self.vaults.get_key(kid)
|
||||||
|
vault_key, vault = looked_up[kid]
|
||||||
if not vault_key or vault_key == server_key:
|
if not vault_key or vault_key == server_key:
|
||||||
continue
|
continue
|
||||||
drm.content_keys[kid] = vault_key
|
drm.content_keys[kid] = vault_key
|
||||||
@@ -4579,10 +4587,10 @@ class dl:
|
|||||||
"""Cache the keys a batch licence or an import put on the tracks.
|
"""Cache the keys a batch licence or an import put on the tracks.
|
||||||
|
|
||||||
Both fill track.drm before any prepare_drm call, so this is the only point where those
|
Both fill track.drm before any prepare_drm call, so this is the only point where those
|
||||||
keys can still reach the local vaults.
|
keys can still reach the local vaults. A content key a local vault already holds is not pushed again.
|
||||||
"""
|
"""
|
||||||
server_vault_keys = getattr(getattr(self, "_remote_service", None), "server_vault_keys", {})
|
server_vault_keys = getattr(getattr(self, "_remote_service", None), "server_vault_keys", {})
|
||||||
keys = {
|
resolved = {
|
||||||
kid: key
|
kid: key
|
||||||
for track in title.tracks
|
for track in title.tracks
|
||||||
for drm in (getattr(track, "drm", None) or [])
|
for drm in (getattr(track, "drm", None) or [])
|
||||||
@@ -4591,6 +4599,7 @@ class dl:
|
|||||||
and server_vault_keys.get(kid) != key
|
and server_vault_keys.get(kid) != key
|
||||||
and self.vaults.sources.get(kid, (None,))[0] != key
|
and self.vaults.sources.get(kid, (None,))[0] != key
|
||||||
}
|
}
|
||||||
|
keys = {kid: key for kid, key in resolved.items() if not self.vaults.held_locally(kid, key)}
|
||||||
if keys:
|
if keys:
|
||||||
self.LICENSE_KEY_CACHE.update(keys)
|
self.LICENSE_KEY_CACHE.update(keys)
|
||||||
self.flush_vault_writes([partial(self.cache_keys_to_vaults, keys)])
|
self.flush_vault_writes([partial(self.cache_keys_to_vaults, keys)])
|
||||||
|
|||||||
@@ -38,7 +38,7 @@ from unshackle.core.config import config
|
|||||||
from unshackle.core.console import console, prompt_user
|
from unshackle.core.console import console, prompt_user
|
||||||
from unshackle.core.constants import AnyTrack
|
from unshackle.core.constants import AnyTrack
|
||||||
from unshackle.core.credential import Credential
|
from unshackle.core.credential import Credential
|
||||||
from unshackle.core.drm import cbcs
|
from unshackle.core.drm import cbcs, real_kids
|
||||||
from unshackle.core.proxies.resolve import is_loopback
|
from unshackle.core.proxies.resolve import is_loopback
|
||||||
from unshackle.core.session import MAX_BACKOFF
|
from unshackle.core.session import MAX_BACKOFF
|
||||||
from unshackle.core.titles import Title_T, Titles_T, remap_titles
|
from unshackle.core.titles import Title_T, Titles_T, remap_titles
|
||||||
@@ -303,7 +303,8 @@ def reconstruct_drm(drm_list: Optional[list]) -> list:
|
|||||||
from unshackle.core.drm import Widevine
|
from unshackle.core.drm import Widevine
|
||||||
|
|
||||||
wv_pssh = WidevinePSSH(pssh_str)
|
wv_pssh = WidevinePSSH(pssh_str)
|
||||||
result.append(Widevine(pssh=wv_pssh))
|
kids = drm_info.get("kids")
|
||||||
|
result.append(Widevine(pssh=wv_pssh, kid=kids[0] if kids and not wv_pssh.key_ids else None))
|
||||||
elif drm_type == "playready":
|
elif drm_type == "playready":
|
||||||
import base64 as b64
|
import base64 as b64
|
||||||
|
|
||||||
@@ -1542,9 +1543,9 @@ class RemoteService:
|
|||||||
self.log.warning(f"Server CDM returned no content keys for track {track.id}")
|
self.log.warning(f"Server CDM returned no content keys for track {track.id}")
|
||||||
continue
|
continue
|
||||||
|
|
||||||
kid_list = list(track_keys.keys())
|
named = {kid.hex for drm in track.drm or [] for kid in real_kids(drm)} & track_keys.keys()
|
||||||
track_drm_type = drm_types_by_track.get(str(track.id), server_drm_type)
|
track_drm_type = drm_types_by_track.get(str(track.id), server_drm_type)
|
||||||
drm_obj = self.create_drm_stub(track_drm_type, kid_list)
|
drm_obj = self.create_drm_stub(track_drm_type, sorted(named) or list(track_keys))
|
||||||
for kid_hex, key_hex in track_keys.items():
|
for kid_hex, key_hex in track_keys.items():
|
||||||
drm_obj.content_keys[UUID(hex=kid_hex)] = key_hex
|
drm_obj.content_keys[UUID(hex=kid_hex)] = key_hex
|
||||||
track.drm = [drm_obj]
|
track.drm = [drm_obj]
|
||||||
|
|||||||
@@ -2,6 +2,46 @@ from abc import ABCMeta, abstractmethod
|
|||||||
from typing import Iterator, Optional, Union
|
from typing import Iterator, Optional, Union
|
||||||
from uuid import UUID
|
from uuid import UUID
|
||||||
|
|
||||||
|
from requests import ConnectionError as RequestsConnectionError
|
||||||
|
from requests import Session, Timeout
|
||||||
|
from requests.adapters import HTTPAdapter
|
||||||
|
from urllib3.util.retry import Retry
|
||||||
|
|
||||||
|
from unshackle.core import __version__
|
||||||
|
from unshackle.core.session import POOL_MAX_SIZE
|
||||||
|
|
||||||
|
|
||||||
|
class RateLimited(Exception):
|
||||||
|
"""The vault still answered 429 after the retries."""
|
||||||
|
|
||||||
|
|
||||||
|
UNREACHABLE = (RequestsConnectionError, Timeout, ConnectionError, TimeoutError)
|
||||||
|
|
||||||
|
|
||||||
|
def vault_session(timeout: float) -> Session:
|
||||||
|
"""An HTTP session for a remote vault that retries a 429 and nothing else.
|
||||||
|
|
||||||
|
A 429 means the request was not processed, so a POST is safe to send again. A connect or read
|
||||||
|
failure is not retried: a lookup waits on every remote vault, so a retry multiplies the timeout.
|
||||||
|
"""
|
||||||
|
retry = Retry(
|
||||||
|
total=3,
|
||||||
|
connect=0,
|
||||||
|
read=0,
|
||||||
|
other=0,
|
||||||
|
backoff_factor=0.5,
|
||||||
|
status_forcelist=[429],
|
||||||
|
allowed_methods=None,
|
||||||
|
retry_after_max=int(timeout),
|
||||||
|
raise_on_status=False,
|
||||||
|
)
|
||||||
|
session = Session()
|
||||||
|
adapter = HTTPAdapter(pool_connections=POOL_MAX_SIZE, pool_maxsize=POOL_MAX_SIZE, max_retries=retry)
|
||||||
|
session.mount("https://", adapter)
|
||||||
|
session.mount("http://", adapter)
|
||||||
|
session.headers.update({"User-Agent": f"unshackle v{__version__}"})
|
||||||
|
return session
|
||||||
|
|
||||||
|
|
||||||
class Vault(metaclass=ABCMeta):
|
class Vault(metaclass=ABCMeta):
|
||||||
# a vault on this machine; Vaults.get_key tries these first, before the network ones
|
# a vault on this machine; Vaults.get_key tries these first, before the network ones
|
||||||
|
|||||||
@@ -8,7 +8,7 @@ from uuid import UUID
|
|||||||
|
|
||||||
from unshackle.core.config import config
|
from unshackle.core.config import config
|
||||||
from unshackle.core.utilities import get_debug_logger, import_module_by_path
|
from unshackle.core.utilities import get_debug_logger, import_module_by_path
|
||||||
from unshackle.core.vault import Vault
|
from unshackle.core.vault import UNREACHABLE, RateLimited, Vault
|
||||||
|
|
||||||
log = logging.getLogger(__name__)
|
log = logging.getLogger(__name__)
|
||||||
|
|
||||||
@@ -29,6 +29,7 @@ class Vaults:
|
|||||||
self.flagged: set[tuple[Union[UUID, str], str]] = set()
|
self.flagged: set[tuple[Union[UUID, str], str]] = set()
|
||||||
self.candidates: dict[Union[UUID, str], list[tuple[str, Vault]]] = {}
|
self.candidates: dict[Union[UUID, str], list[tuple[str, Vault]]] = {}
|
||||||
self.warned: set[tuple[str, Union[UUID, str], str]] = set()
|
self.warned: set[tuple[str, Union[UUID, str], str]] = set()
|
||||||
|
self.unreachable: set[Vault] = set()
|
||||||
|
|
||||||
def __iter__(self) -> Iterator[Vault]:
|
def __iter__(self) -> Iterator[Vault]:
|
||||||
return iter(self.vaults)
|
return iter(self.vaults)
|
||||||
@@ -70,7 +71,7 @@ class Vaults:
|
|||||||
and the first one with the content key wins.
|
and the first one with the content key wins.
|
||||||
"""
|
"""
|
||||||
local = [v for v in self.vaults if v.local]
|
local = [v for v in self.vaults if v.local]
|
||||||
remote = [v for v in self.vaults if not v.local]
|
remote = [v for v in self.live() if not v.local]
|
||||||
for vault in local:
|
for vault in local:
|
||||||
key = self.query_vault(vault, kid)
|
key = self.query_vault(vault, kid)
|
||||||
if key:
|
if key:
|
||||||
@@ -123,6 +124,7 @@ class Vaults:
|
|||||||
return None
|
return None
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
log.warning(f"Failed to get key from Vault '{vault.name}': {e}")
|
log.warning(f"Failed to get key from Vault '{vault.name}': {e}")
|
||||||
|
self.mark_unreachable(vault, e)
|
||||||
if dl:
|
if dl:
|
||||||
dl.log_vault_query(
|
dl.log_vault_query(
|
||||||
vault.name,
|
vault.name,
|
||||||
@@ -160,7 +162,7 @@ class Vaults:
|
|||||||
"""
|
"""
|
||||||
self.flagged.add((kid, key))
|
self.flagged.add((kid, key))
|
||||||
source = self.sources.pop(kid, None)
|
source = self.sources.pop(kid, None)
|
||||||
for vault in self.vaults:
|
for vault in self.live():
|
||||||
if not vault.local and (source is None or vault is not source[1]):
|
if not vault.local and (source is None or vault is not source[1]):
|
||||||
continue
|
continue
|
||||||
try:
|
try:
|
||||||
@@ -171,7 +173,7 @@ class Vaults:
|
|||||||
def unflag_bad_key(self, kid: Union[UUID, str], key: str) -> None:
|
def unflag_bad_key(self, kid: Union[UUID, str], key: str) -> None:
|
||||||
"""Remove a flag from every Vault."""
|
"""Remove a flag from every Vault."""
|
||||||
self.flagged.discard((kid, key))
|
self.flagged.discard((kid, key))
|
||||||
for vault in self.vaults:
|
for vault in self.live():
|
||||||
try:
|
try:
|
||||||
vault.unflag_bad_key(kid, key)
|
vault.unflag_bad_key(kid, key)
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
@@ -183,7 +185,7 @@ class Vaults:
|
|||||||
This method pushes to all Vaults at the same time, so one unreachable Vault costs only its
|
This method pushes to all Vaults at the same time, so one unreachable Vault costs only its
|
||||||
own timeout and does not delay the others.
|
own timeout and does not delay the others.
|
||||||
"""
|
"""
|
||||||
vaults = [vault for vault in self.vaults if vault != excluding and not vault.no_push]
|
vaults = [v for v in self.live() if v != excluding and not v.no_push]
|
||||||
if not vaults:
|
if not vaults:
|
||||||
return 0
|
return 0
|
||||||
with ThreadPoolExecutor(len(vaults), thread_name_prefix="vault-push") as pool:
|
with ThreadPoolExecutor(len(vaults), thread_name_prefix="vault-push") as pool:
|
||||||
@@ -201,6 +203,7 @@ class Vaults:
|
|||||||
return 0
|
return 0
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
log.warning(f"Failed to add key to Vault '{vault.name}': {e}")
|
log.warning(f"Failed to add key to Vault '{vault.name}': {e}")
|
||||||
|
self.mark_unreachable(vault, e)
|
||||||
if dl:
|
if dl:
|
||||||
dl.log_vault_query(vault.name, "add_key", kid=str(kid), success=False, error=e)
|
dl.log_vault_query(vault.name, "add_key", kid=str(kid), success=False, error=e)
|
||||||
return 0
|
return 0
|
||||||
@@ -214,7 +217,7 @@ class Vaults:
|
|||||||
This method pushes to all Vaults at the same time, so one unreachable Vault costs only its
|
This method pushes to all Vaults at the same time, so one unreachable Vault costs only its
|
||||||
own timeout and does not delay the others.
|
own timeout and does not delay the others.
|
||||||
"""
|
"""
|
||||||
vaults = [vault for vault in self.vaults if not vault.no_push]
|
vaults = [vault for vault in self.live() if not vault.no_push]
|
||||||
if not vaults:
|
if not vaults:
|
||||||
return 0
|
return 0
|
||||||
with ThreadPoolExecutor(len(vaults), thread_name_prefix="vault-push") as pool:
|
with ThreadPoolExecutor(len(vaults), thread_name_prefix="vault-push") as pool:
|
||||||
@@ -229,7 +232,23 @@ class Vaults:
|
|||||||
return 0
|
return 0
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
log.warning(f"Failed to add keys to Vault '{vault.name}': {e}")
|
log.warning(f"Failed to add keys to Vault '{vault.name}': {e}")
|
||||||
|
self.mark_unreachable(vault, e)
|
||||||
return 0
|
return 0
|
||||||
|
|
||||||
|
def live(self) -> list[Vault]:
|
||||||
|
"""The vaults this run still uses."""
|
||||||
|
return [vault for vault in self.vaults if vault not in self.unreachable]
|
||||||
|
|
||||||
|
def held_locally(self, kid: Union[UUID, str], key: str) -> bool:
|
||||||
|
"""True when a local vault already holds this KID:KEY."""
|
||||||
|
return any(self.query_vault(vault, kid) == key for vault in self.vaults if vault.local)
|
||||||
|
|
||||||
|
def mark_unreachable(self, vault: Vault, error: Exception) -> None:
|
||||||
|
"""Skip a vault that is down or still rate limited for the rest of the run, so it costs one
|
||||||
|
timeout, not one per KID."""
|
||||||
|
if isinstance(error, (RateLimited, *UNREACHABLE)) and vault not in self.unreachable:
|
||||||
|
self.unreachable.add(vault)
|
||||||
|
log.warning(f"Vault '{vault.name}' is unavailable, so this run does not use it again")
|
||||||
|
|
||||||
|
|
||||||
__all__ = ("Vaults",)
|
__all__ = ("Vaults",)
|
||||||
|
|||||||
+10
-13
@@ -1,12 +1,9 @@
|
|||||||
from typing import Iterator, Optional, Union
|
from typing import Iterator, Optional, Union
|
||||||
from uuid import UUID
|
from uuid import UUID
|
||||||
|
|
||||||
from requests import Response, Session
|
from requests import HTTPError, Response
|
||||||
from requests.adapters import HTTPAdapter
|
|
||||||
|
|
||||||
from unshackle.core import __version__
|
from unshackle.core.vault import RateLimited, Vault, vault_session
|
||||||
from unshackle.core.session import POOL_MAX_SIZE
|
|
||||||
from unshackle.core.vault import Vault
|
|
||||||
|
|
||||||
|
|
||||||
class API(Vault):
|
class API(Vault):
|
||||||
@@ -24,11 +21,7 @@ class API(Vault):
|
|||||||
super().__init__(name, no_push)
|
super().__init__(name, no_push)
|
||||||
self.uri = uri.rstrip("/")
|
self.uri = uri.rstrip("/")
|
||||||
self.timeout = timeout
|
self.timeout = timeout
|
||||||
self.session = Session()
|
self.session = vault_session(timeout)
|
||||||
adapter = HTTPAdapter(pool_connections=POOL_MAX_SIZE, pool_maxsize=POOL_MAX_SIZE)
|
|
||||||
self.session.mount("https://", adapter)
|
|
||||||
self.session.mount("http://", adapter)
|
|
||||||
self.session.headers.update({"User-Agent": f"unshackle v{__version__}"})
|
|
||||||
self.session.headers.update({"Authorization": f"Bearer {token}"})
|
self.session.headers.update({"Authorization": f"Bearer {token}"})
|
||||||
|
|
||||||
if headers:
|
if headers:
|
||||||
@@ -44,6 +37,8 @@ class API(Vault):
|
|||||||
timeout=self.timeout,
|
timeout=self.timeout,
|
||||||
allow_redirects=False,
|
allow_redirects=False,
|
||||||
)
|
)
|
||||||
|
if response.status_code == 429:
|
||||||
|
raise Exceptions.TooManyRequests(f"Vault {self.name} rate limited the lookup")
|
||||||
response.raise_for_status()
|
response.raise_for_status()
|
||||||
data = response.json()
|
data = response.json()
|
||||||
|
|
||||||
@@ -176,10 +171,12 @@ class API(Vault):
|
|||||||
batch_size = 1
|
batch_size = 1
|
||||||
continue
|
continue
|
||||||
|
|
||||||
|
if response.status_code == 429:
|
||||||
|
raise Exceptions.TooManyRequests(f"Vault {self.name} rate limited the push")
|
||||||
response.raise_for_status()
|
response.raise_for_status()
|
||||||
data = response.json()
|
data = response.json()
|
||||||
except Exception:
|
except (ValueError, HTTPError):
|
||||||
# JSON decode error or connection issue - try smaller batch
|
# an HTTP error or a broken JSON reply - try smaller batch
|
||||||
if batch_size > 1:
|
if batch_size > 1:
|
||||||
if batch_size > 100:
|
if batch_size > 100:
|
||||||
batch_size -= 100
|
batch_size -= 100
|
||||||
@@ -281,7 +278,7 @@ class Exceptions:
|
|||||||
class AuthRejected(Exception):
|
class AuthRejected(Exception):
|
||||||
"""Authentication Error Occurred, is your token valid? Do you have permission to make this call?"""
|
"""Authentication Error Occurred, is your token valid? Do you have permission to make this call?"""
|
||||||
|
|
||||||
class TooManyRequests(Exception):
|
class TooManyRequests(RateLimited):
|
||||||
"""Rate Limited. Sent too many requests in a given amount of time."""
|
"""Rate Limited. Sent too many requests in a given amount of time."""
|
||||||
|
|
||||||
class PageInvalid(Exception):
|
class PageInvalid(Exception):
|
||||||
|
|||||||
@@ -3,12 +3,7 @@ from enum import Enum
|
|||||||
from typing import Iterator, Optional, Union
|
from typing import Iterator, Optional, Union
|
||||||
from uuid import UUID
|
from uuid import UUID
|
||||||
|
|
||||||
from requests import Session
|
from unshackle.core.vault import RateLimited, Vault, vault_session
|
||||||
from requests.adapters import HTTPAdapter
|
|
||||||
|
|
||||||
from unshackle.core import __version__
|
|
||||||
from unshackle.core.session import POOL_MAX_SIZE
|
|
||||||
from unshackle.core.vault import Vault
|
|
||||||
|
|
||||||
|
|
||||||
class InsertResult(Enum):
|
class InsertResult(Enum):
|
||||||
@@ -61,11 +56,7 @@ class HTTP(Vault):
|
|||||||
self.api_mode = api_mode.lower()
|
self.api_mode = api_mode.lower()
|
||||||
self.batch_insert = True
|
self.batch_insert = True
|
||||||
self.current_title = None
|
self.current_title = None
|
||||||
self.session = Session()
|
self.session = vault_session(timeout)
|
||||||
adapter = HTTPAdapter(pool_connections=POOL_MAX_SIZE, pool_maxsize=POOL_MAX_SIZE)
|
|
||||||
self.session.mount("https://", adapter)
|
|
||||||
self.session.mount("http://", adapter)
|
|
||||||
self.session.headers.update({"User-Agent": f"unshackle v{__version__}"})
|
|
||||||
self.api_session_id = None
|
self.api_session_id = None
|
||||||
|
|
||||||
if self.api_mode == "decrypt_labs":
|
if self.api_mode == "decrypt_labs":
|
||||||
@@ -97,6 +88,9 @@ class HTTP(Vault):
|
|||||||
if r.status_code == 404:
|
if r.status_code == 404:
|
||||||
return {"status": "not_found"}
|
return {"status": "not_found"}
|
||||||
|
|
||||||
|
if r.status_code == 429:
|
||||||
|
raise RateLimited(f"Vault {self.name} rate limited the request")
|
||||||
|
|
||||||
if not r.ok:
|
if not r.ok:
|
||||||
try:
|
try:
|
||||||
detail = r.json().get("message") or r.reason.title()
|
detail = r.json().get("message") or r.reason.title()
|
||||||
@@ -124,11 +118,8 @@ class HTTP(Vault):
|
|||||||
kid = kid.hex
|
kid = kid.hex
|
||||||
|
|
||||||
if self.api_mode == "decrypt_labs":
|
if self.api_mode == "decrypt_labs":
|
||||||
|
response = self.session.post(self.url, json={"service": service.lower(), "kid": kid}, timeout=self.timeout)
|
||||||
try:
|
try:
|
||||||
request_payload = {"service": service.lower(), "kid": kid}
|
|
||||||
|
|
||||||
response = self.session.post(self.url, json=request_payload, timeout=self.timeout)
|
|
||||||
|
|
||||||
if not response.ok:
|
if not response.ok:
|
||||||
return None
|
return None
|
||||||
|
|
||||||
@@ -186,7 +177,7 @@ class HTTP(Vault):
|
|||||||
elif isinstance(key_entry, dict):
|
elif isinstance(key_entry, dict):
|
||||||
if key_entry.get("kid") == kid:
|
if key_entry.get("kid") == kid:
|
||||||
return key_entry.get("key")
|
return key_entry.get("key")
|
||||||
except Exception as e:
|
except ValueError as e:
|
||||||
print(f"Failed to get key ({e.__class__.__name__}: {e})")
|
print(f"Failed to get key ({e.__class__.__name__}: {e})")
|
||||||
return None
|
return None
|
||||||
return None
|
return None
|
||||||
@@ -300,8 +291,6 @@ class HTTP(Vault):
|
|||||||
except ValueError as e:
|
except ValueError as e:
|
||||||
if "unknown" in str(e).lower() and "method" in str(e).lower():
|
if "unknown" in str(e).lower() and "method" in str(e).lower():
|
||||||
self.batch_insert = False
|
self.batch_insert = False
|
||||||
except Exception:
|
|
||||||
pass
|
|
||||||
for kid, key in processed_kid_keys.items():
|
for kid, key in processed_kid_keys.items():
|
||||||
try:
|
try:
|
||||||
response = self.request(
|
response = self.request(
|
||||||
@@ -317,7 +306,7 @@ class HTTP(Vault):
|
|||||||
continue
|
continue
|
||||||
if response.get("inserted", False):
|
if response.get("inserted", False):
|
||||||
inserted_count += 1
|
inserted_count += 1
|
||||||
except Exception as e:
|
except ValueError as e:
|
||||||
failures.append(e)
|
failures.append(e)
|
||||||
if failures and len(failures) == len(processed_kid_keys):
|
if failures and len(failures) == len(processed_kid_keys):
|
||||||
raise failures[-1]
|
raise failures[-1]
|
||||||
|
|||||||
Reference in New Issue
Block a user