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
|
||||
constructor accepts a `timeout` parameter (a per-vault `timeout` still wins). The `MySQL` backend
|
||||
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.dl import dl
|
||||
from unshackle.core.vault import RateLimited, vault_session
|
||||
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
|
||||
|
||||
KID = UUID("11111111-2222-3333-4444-555555555555")
|
||||
SIBLING = UUID("99999999-2222-3333-4444-555555555555")
|
||||
BAD = "aa" * 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.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))
|
||||
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])))
|
||||
|
||||
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):
|
||||
return 1
|
||||
|
||||
def held_locally(self, kid, key):
|
||||
return False
|
||||
|
||||
def get_key(self, kid):
|
||||
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.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.export_name import ExportNamer, move_export, plural, season_episodes, title_label
|
||||
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
|
||||
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
|
||||
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", {})
|
||||
looked_up: dict[UUID, tuple[Optional[str], Any]] = {}
|
||||
for track in title.tracks:
|
||||
for drm in getattr(track, "drm", None) or []:
|
||||
for kid, server_key in list(getattr(drm, "content_keys", {}).items()):
|
||||
vault_key, vault = self.vaults.get_key(kid)
|
||||
content_keys = getattr(drm, "content_keys", {})
|
||||
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:
|
||||
continue
|
||||
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.
|
||||
|
||||
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", {})
|
||||
keys = {
|
||||
resolved = {
|
||||
kid: key
|
||||
for track in title.tracks
|
||||
for drm in (getattr(track, "drm", None) or [])
|
||||
@@ -4591,6 +4599,7 @@ class dl:
|
||||
and server_vault_keys.get(kid) != 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:
|
||||
self.LICENSE_KEY_CACHE.update(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.constants import AnyTrack
|
||||
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.session import MAX_BACKOFF
|
||||
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
|
||||
|
||||
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":
|
||||
import base64 as b64
|
||||
|
||||
@@ -1542,9 +1543,9 @@ class RemoteService:
|
||||
self.log.warning(f"Server CDM returned no content keys for track {track.id}")
|
||||
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)
|
||||
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():
|
||||
drm_obj.content_keys[UUID(hex=kid_hex)] = key_hex
|
||||
track.drm = [drm_obj]
|
||||
|
||||
@@ -2,6 +2,46 @@ from abc import ABCMeta, abstractmethod
|
||||
from typing import Iterator, Optional, Union
|
||||
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):
|
||||
# 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.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__)
|
||||
|
||||
@@ -29,6 +29,7 @@ class Vaults:
|
||||
self.flagged: set[tuple[Union[UUID, str], str]] = set()
|
||||
self.candidates: dict[Union[UUID, str], list[tuple[str, Vault]]] = {}
|
||||
self.warned: set[tuple[str, Union[UUID, str], str]] = set()
|
||||
self.unreachable: set[Vault] = set()
|
||||
|
||||
def __iter__(self) -> Iterator[Vault]:
|
||||
return iter(self.vaults)
|
||||
@@ -70,7 +71,7 @@ class Vaults:
|
||||
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]
|
||||
remote = [v for v in self.live() if not v.local]
|
||||
for vault in local:
|
||||
key = self.query_vault(vault, kid)
|
||||
if key:
|
||||
@@ -123,6 +124,7 @@ class Vaults:
|
||||
return None
|
||||
except Exception as e:
|
||||
log.warning(f"Failed to get key from Vault '{vault.name}': {e}")
|
||||
self.mark_unreachable(vault, e)
|
||||
if dl:
|
||||
dl.log_vault_query(
|
||||
vault.name,
|
||||
@@ -160,7 +162,7 @@ class Vaults:
|
||||
"""
|
||||
self.flagged.add((kid, key))
|
||||
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]):
|
||||
continue
|
||||
try:
|
||||
@@ -171,7 +173,7 @@ class Vaults:
|
||||
def unflag_bad_key(self, kid: Union[UUID, str], key: str) -> None:
|
||||
"""Remove a flag from every Vault."""
|
||||
self.flagged.discard((kid, key))
|
||||
for vault in self.vaults:
|
||||
for vault in self.live():
|
||||
try:
|
||||
vault.unflag_bad_key(kid, key)
|
||||
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
|
||||
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:
|
||||
return 0
|
||||
with ThreadPoolExecutor(len(vaults), thread_name_prefix="vault-push") as pool:
|
||||
@@ -201,6 +203,7 @@ class Vaults:
|
||||
return 0
|
||||
except Exception as e:
|
||||
log.warning(f"Failed to add key to Vault '{vault.name}': {e}")
|
||||
self.mark_unreachable(vault, e)
|
||||
if dl:
|
||||
dl.log_vault_query(vault.name, "add_key", kid=str(kid), success=False, error=e)
|
||||
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
|
||||
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:
|
||||
return 0
|
||||
with ThreadPoolExecutor(len(vaults), thread_name_prefix="vault-push") as pool:
|
||||
@@ -229,7 +232,23 @@ class Vaults:
|
||||
return 0
|
||||
except Exception as e:
|
||||
log.warning(f"Failed to add keys to Vault '{vault.name}': {e}")
|
||||
self.mark_unreachable(vault, e)
|
||||
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",)
|
||||
|
||||
+10
-13
@@ -1,12 +1,9 @@
|
||||
from typing import Iterator, Optional, Union
|
||||
from uuid import UUID
|
||||
|
||||
from requests import Response, Session
|
||||
from requests.adapters import HTTPAdapter
|
||||
from requests import HTTPError, Response
|
||||
|
||||
from unshackle.core import __version__
|
||||
from unshackle.core.session import POOL_MAX_SIZE
|
||||
from unshackle.core.vault import Vault
|
||||
from unshackle.core.vault import RateLimited, Vault, vault_session
|
||||
|
||||
|
||||
class API(Vault):
|
||||
@@ -24,11 +21,7 @@ class API(Vault):
|
||||
super().__init__(name, no_push)
|
||||
self.uri = uri.rstrip("/")
|
||||
self.timeout = timeout
|
||||
self.session = Session()
|
||||
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 = vault_session(timeout)
|
||||
self.session.headers.update({"Authorization": f"Bearer {token}"})
|
||||
|
||||
if headers:
|
||||
@@ -44,6 +37,8 @@ class API(Vault):
|
||||
timeout=self.timeout,
|
||||
allow_redirects=False,
|
||||
)
|
||||
if response.status_code == 429:
|
||||
raise Exceptions.TooManyRequests(f"Vault {self.name} rate limited the lookup")
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
|
||||
@@ -176,10 +171,12 @@ class API(Vault):
|
||||
batch_size = 1
|
||||
continue
|
||||
|
||||
if response.status_code == 429:
|
||||
raise Exceptions.TooManyRequests(f"Vault {self.name} rate limited the push")
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
except Exception:
|
||||
# JSON decode error or connection issue - try smaller batch
|
||||
except (ValueError, HTTPError):
|
||||
# an HTTP error or a broken JSON reply - try smaller batch
|
||||
if batch_size > 1:
|
||||
if batch_size > 100:
|
||||
batch_size -= 100
|
||||
@@ -281,7 +278,7 @@ class Exceptions:
|
||||
class AuthRejected(Exception):
|
||||
"""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."""
|
||||
|
||||
class PageInvalid(Exception):
|
||||
|
||||
@@ -3,12 +3,7 @@ from enum import Enum
|
||||
from typing import Iterator, Optional, Union
|
||||
from uuid import UUID
|
||||
|
||||
from requests import 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
|
||||
from unshackle.core.vault import RateLimited, Vault, vault_session
|
||||
|
||||
|
||||
class InsertResult(Enum):
|
||||
@@ -61,11 +56,7 @@ class HTTP(Vault):
|
||||
self.api_mode = api_mode.lower()
|
||||
self.batch_insert = True
|
||||
self.current_title = None
|
||||
self.session = Session()
|
||||
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 = vault_session(timeout)
|
||||
self.api_session_id = None
|
||||
|
||||
if self.api_mode == "decrypt_labs":
|
||||
@@ -97,6 +88,9 @@ class HTTP(Vault):
|
||||
if r.status_code == 404:
|
||||
return {"status": "not_found"}
|
||||
|
||||
if r.status_code == 429:
|
||||
raise RateLimited(f"Vault {self.name} rate limited the request")
|
||||
|
||||
if not r.ok:
|
||||
try:
|
||||
detail = r.json().get("message") or r.reason.title()
|
||||
@@ -124,11 +118,8 @@ class HTTP(Vault):
|
||||
kid = kid.hex
|
||||
|
||||
if self.api_mode == "decrypt_labs":
|
||||
response = self.session.post(self.url, json={"service": service.lower(), "kid": kid}, timeout=self.timeout)
|
||||
try:
|
||||
request_payload = {"service": service.lower(), "kid": kid}
|
||||
|
||||
response = self.session.post(self.url, json=request_payload, timeout=self.timeout)
|
||||
|
||||
if not response.ok:
|
||||
return None
|
||||
|
||||
@@ -186,7 +177,7 @@ class HTTP(Vault):
|
||||
elif isinstance(key_entry, dict):
|
||||
if key_entry.get("kid") == kid:
|
||||
return key_entry.get("key")
|
||||
except Exception as e:
|
||||
except ValueError as e:
|
||||
print(f"Failed to get key ({e.__class__.__name__}: {e})")
|
||||
return None
|
||||
return None
|
||||
@@ -300,8 +291,6 @@ class HTTP(Vault):
|
||||
except ValueError as e:
|
||||
if "unknown" in str(e).lower() and "method" in str(e).lower():
|
||||
self.batch_insert = False
|
||||
except Exception:
|
||||
pass
|
||||
for kid, key in processed_kid_keys.items():
|
||||
try:
|
||||
response = self.request(
|
||||
@@ -317,7 +306,7 @@ class HTTP(Vault):
|
||||
continue
|
||||
if response.get("inserted", False):
|
||||
inserted_count += 1
|
||||
except Exception as e:
|
||||
except ValueError as e:
|
||||
failures.append(e)
|
||||
if failures and len(failures) == len(processed_kid_keys):
|
||||
raise failures[-1]
|
||||
|
||||
Reference in New Issue
Block a user