diff --git a/unshackle/core/config.py b/unshackle/core/config.py index 44aad5e4..8be0e54b 100644 --- a/unshackle/core/config.py +++ b/unshackle/core/config.py @@ -75,6 +75,7 @@ class Config: self.audio: dict = kwargs.get("audio") or {} self.headers: dict = kwargs.get("headers") or {} self.key_vaults: list[dict[str, Any]] = kwargs.get("key_vaults", []) + self.vault_timeout: float = kwargs.get("vault_timeout", 10.0) self.muxing: dict = kwargs.get("muxing") or {} self.proxy_providers: dict = kwargs.get("proxy_providers") or {} self.remote_services: dict = kwargs.get("remote_services") or {} diff --git a/unshackle/core/vaults.py b/unshackle/core/vaults.py index 3db74eca..538eec3e 100644 --- a/unshackle/core/vaults.py +++ b/unshackle/core/vaults.py @@ -1,3 +1,4 @@ +import inspect import logging import time from typing import Any, Iterator, Optional, Union @@ -34,6 +35,9 @@ class Vaults: module = _MODULES.get(type_) if not module: raise ValueError(f"Unable to find vault command by the name '{type_}'.") + # only pass the global default to vaults that declare a timeout param (per-vault config still wins) + if "timeout" not in kwargs and "timeout" in inspect.signature(module).parameters: + kwargs["timeout"] = config.vault_timeout try: vault = module(**kwargs) self.vaults.append(vault) diff --git a/unshackle/unshackle-example.yaml b/unshackle/unshackle-example.yaml index bcb47c19..c1ccf367 100644 --- a/unshackle/unshackle-example.yaml +++ b/unshackle/unshackle-example.yaml @@ -394,6 +394,11 @@ remote_cdm: # Key Vaults store your obtained Content Encryption Keys (CEKs) # Use 'no_push: true' to prevent a vault from receiving pushed keys # while still allowing it to provide keys when requested + +# Global network timeout (seconds) for remote vaults (API, HTTP). Defaults to 10. +# Override per-vault with a 'timeout:' key on that vault entry. +# vault_timeout: 10.0 + key_vaults: - type: SQLite name: Local @@ -403,6 +408,7 @@ key_vaults: # name: "Remote Vault" # uri: "https://key-vault.example.com" # token: "secret_token" + # timeout: 5.0 # overrides the global vault_timeout for this vault # no_push: true # This vault will only provide keys, not receive them # - type: MySQL # name: "MySQL Vault" diff --git a/unshackle/vaults/API.py b/unshackle/vaults/API.py index dad9607b..257313d3 100644 --- a/unshackle/vaults/API.py +++ b/unshackle/vaults/API.py @@ -10,9 +10,10 @@ from unshackle.core.vault import Vault class API(Vault): """Key Vault using a simple RESTful HTTP API call.""" - def __init__(self, name: str, uri: str, token: str, no_push: bool = False): + def __init__(self, name: str, uri: str, token: str, no_push: bool = False, timeout: float = 10.0): super().__init__(name, no_push) self.uri = uri.rstrip("/") + self.timeout = timeout self.session = Session() self.session.headers.update({"User-Agent": f"unshackle v{__version__}"}) self.session.headers.update({"Authorization": f"Bearer {token}"}) @@ -22,7 +23,7 @@ class API(Vault): kid = kid.hex data = self.session.get( - url=f"{self.uri}/{service.lower()}/{kid}", headers={"Accept": "application/json"} + url=f"{self.uri}/{service.lower()}/{kid}", headers={"Accept": "application/json"}, timeout=self.timeout ).json() code = int(data.get("code", 0)) @@ -55,6 +56,7 @@ class API(Vault): url=f"{self.uri}/{service.lower()}", params={"page": page, "total": 10}, headers={"Accept": "application/json"}, + timeout=self.timeout, ).json() code = int(data.get("code", 0)) @@ -89,7 +91,10 @@ class API(Vault): kid = kid.hex data = self.session.post( - url=f"{self.uri}/{service.lower()}/{kid}", json={"content_key": key}, headers={"Accept": "application/json"} + url=f"{self.uri}/{service.lower()}/{kid}", + json={"content_key": key}, + headers={"Accept": "application/json"}, + timeout=self.timeout, ).json() code = int(data.get("code", 0)) @@ -135,6 +140,7 @@ class API(Vault): url=f"{self.uri}/{service.lower()}", json={"content_keys": batch_keys}, headers={"Accept": "application/json"}, + timeout=self.timeout, ) # Check for HTTP errors that suggest batch is too large @@ -181,7 +187,7 @@ class API(Vault): return total_added def get_services(self) -> Iterator[str]: - data = self.session.post(url=self.uri, headers={"Accept": "application/json"}).json() + data = self.session.post(url=self.uri, headers={"Accept": "application/json"}, timeout=self.timeout).json() code = int(data.get("code", 0)) message = data.get("message") diff --git a/unshackle/vaults/HTTP.py b/unshackle/vaults/HTTP.py index 65a72f7c..0f5d2ff4 100644 --- a/unshackle/vaults/HTTP.py +++ b/unshackle/vaults/HTTP.py @@ -34,6 +34,7 @@ class HTTP(Vault): username: Optional[str] = None, api_mode: str = "query", no_push: bool = False, + timeout: float = 10.0, ): """ Initialize HTTP Vault. @@ -49,6 +50,7 @@ class HTTP(Vault): """ super().__init__(name, no_push) self.url = host + self.timeout = timeout self.password = api_key or password if not self.password: raise ValueError("Either password or api_key is required") @@ -84,7 +86,7 @@ class HTTP(Vault): "token": self.password, } - r = self.session.post(self.url, json=request_payload) + r = self.session.post(self.url, json=request_payload, timeout=self.timeout) if r.status_code == 404: return {"status": "not_found"} @@ -115,7 +117,7 @@ class HTTP(Vault): try: request_payload = {"service": service.lower(), "kid": kid} - response = self.session.post(self.url, json=request_payload) + response = self.session.post(self.url, json=request_payload, timeout=self.timeout) if not response.ok: return None @@ -182,6 +184,7 @@ class HTTP(Vault): response = self.session.get( self.url, params={"service": service.lower(), "username": self.username, "password": self.password, "kid": kid}, + timeout=self.timeout, ) data = response.json() @@ -200,7 +203,9 @@ class HTTP(Vault): return iter([]) else: # query mode response = self.session.get( - self.url, params={"service": service.lower(), "username": self.username, "password": self.password} + self.url, + params={"service": service.lower(), "username": self.username, "password": self.password}, + timeout=self.timeout, ) data = response.json() @@ -250,6 +255,7 @@ class HTTP(Vault): "key": key, "title": title, }, + timeout=self.timeout, ) data = response.json() @@ -301,6 +307,7 @@ class HTTP(Vault): "key": key, "title": title, }, + timeout=self.timeout, ) data = response.json() @@ -323,7 +330,9 @@ class HTTP(Vault): return iter([]) else: # query mode response = self.session.get( - self.url, params={"username": self.username, "password": self.password, "list_services": True} + self.url, + params={"username": self.username, "password": self.password, "list_services": True}, + timeout=self.timeout, ) data = response.json() @@ -395,6 +404,7 @@ class HTTP(Vault): "key": key, "title": title, }, + timeout=self.timeout, ) try: