diff --git a/packages/envied/src/envied/commands/dl.py b/packages/envied/src/envied/commands/dl.py index fc30329..2fdce5b 100644 --- a/packages/envied/src/envied/commands/dl.py +++ b/packages/envied/src/envied/commands/dl.py @@ -226,12 +226,6 @@ class dl: default=False, help="Skip downloading, only list available titles that would have been downloaded.", ) - @click.option( - "--select-titles", - is_flag=True, - default=False, - help="Interactively select downloads from a list. Only use with Series to select Episodes", - ) @click.option( "--skip-dl", is_flag=True, default=False, help="Skip downloading while still retrieving the decryption keys." ) @@ -255,7 +249,17 @@ class dl: ) @click.option("--downloads", type=int, default=1, help="Amount of tracks to download concurrently.") @click.option("--no-cache", "no_cache", is_flag=True, default=False, help="Bypass title cache for this download.") - @click.option("--reset-cache", "reset_cache", is_flag=True, default=False, help="Clear title cache before fetching.") + @click.option( + "--reset-cache", "reset_cache", is_flag=True, default=False, help="Clear title cache before fetching." + ) + + @click.option( + "--select-titles", + is_flag=True, + default=False, + help="Interactively select downloads from a list. Only use with Series to select Episodes", + ) + @click.pass_context def cli(ctx: click.Context, **kwargs: Any) -> dl: return dl(ctx, **kwargs) @@ -301,20 +305,8 @@ class dl: if getattr(config, "downloader_map", None): config.downloader = config.downloader_map.get(self.service, config.downloader) - with console.status("Loading DRM CDM...", spinner="dots"): - try: - self.cdm = self.get_cdm(self.service, self.profile) - except ValueError as e: - self.log.error(f"Failed to load CDM, {e}") - sys.exit(1) - - if self.cdm: - if hasattr(self.cdm, "device_type") and self.cdm.device_type.name in ["ANDROID", "CHROME"]: - self.log.info(f"Loaded Widevine CDM: {self.cdm.system_id} (L{self.cdm.security_level})") - else: - self.log.info( - f"Loaded PlayReady CDM: {self.cdm.certificate_chain.get_name()} (L{self.cdm.security_level})" - ) + if getattr(config, "decryption_map", None): + config.decryption = config.decryption_map.get(self.service, config.decryption) with console.status("Loading Key Vaults...", spinner="dots"): self.vaults = Vaults(self.service) @@ -354,6 +346,24 @@ class dl: else: self.log.debug("No vaults are currently active") + with console.status("Loading DRM CDM...", spinner="dots"): + try: + self.cdm = self.get_cdm(self.service, self.profile) + except ValueError as e: + self.log.error(f"Failed to load CDM, {e}") + sys.exit(1) + + if self.cdm: + if isinstance(self.cdm, DecryptLabsRemoteCDM): + drm_type = "PlayReady" if self.cdm.is_playready else "Widevine" + self.log.info(f"Loaded {drm_type} Remote CDM: DecryptLabs (L{self.cdm.security_level})") + elif hasattr(self.cdm, "device_type") and self.cdm.device_type.name in ["ANDROID", "CHROME"]: + self.log.info(f"Loaded Widevine CDM: {self.cdm.system_id} (L{self.cdm.security_level})") + else: + self.log.info( + f"Loaded PlayReady CDM: {self.cdm.certificate_chain.get_name()} (L{self.cdm.security_level})" + ) + self.proxy_providers = [] if no_proxy: ctx.params["proxy"] = None @@ -441,7 +451,6 @@ class dl: slow: bool, list_: bool, list_titles: bool, - select_titles: bool, skip_dl: bool, export: Optional[Path], cdm_only: Optional[bool], @@ -450,6 +459,7 @@ class dl: no_source: bool, workers: Optional[int], downloads: int, + select_titles: bool, *_: Any, **__: Any, ) -> None: @@ -500,8 +510,8 @@ class dl: console.print(Padding(titles.tree(verbose=list_titles), (0, 5))) if list_titles: return - - # modification to enable beaupy module to list titles for download for manual selection + + # modification to enable beaupy module to list titles for download for manual selection # use --select-titles after dl in envied command Config.transient = True @@ -535,8 +545,6 @@ class dl: del titles[i] # end modification - - for i, title in enumerate(titles): if isinstance(title, Episode) and wanted and f"{title.season}x{title.number}" not in wanted: continue @@ -575,7 +583,7 @@ class dl: else: console.print(Padding("Search -> [bright_black]No match found[/]", (0, 5))) - if self.tmdb_id and getattr(self, 'search_source', None) != 'simkl': + if self.tmdb_id and getattr(self, "search_source", None) != "simkl": kind = "tv" if isinstance(title, Episode) else "movie" tags.external_ids(self.tmdb_id, kind) if self.tmdb_year: @@ -913,7 +921,12 @@ class dl: ), licence=partial( service.get_playready_license - if isinstance(self.cdm, PlayReadyCdm) + if ( + isinstance(self.cdm, PlayReadyCdm) + or ( + isinstance(self.cdm, DecryptLabsRemoteCDM) and self.cdm.is_playready + ) + ) and hasattr(service, "get_playready_license") else service.get_widevine_license, title=title, @@ -1045,12 +1058,7 @@ class dl: # Handle DRM decryption BEFORE repacking (must decrypt first!) service_name = service.__class__.__name__.upper() decryption_method = config.decryption_map.get(service_name, config.decryption) - use_mp4decrypt = decryption_method.lower() == "mp4decrypt" - - if use_mp4decrypt: - decrypt_tool = "mp4decrypt" - else: - decrypt_tool = "Shaka Packager" + decrypt_tool = "mp4decrypt" if decryption_method.lower() == "mp4decrypt" else "Shaka Packager" drm_tracks = [track for track in title.tracks if track.drm] if drm_tracks: @@ -1059,7 +1067,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, use_mp4decrypt=use_mp4decrypt) + drm.decrypt(track.path) has_decrypted = True events.emit(events.Types.TRACK_REPACKED, track=track) else: @@ -1245,10 +1253,22 @@ class dl: if not drm: return - if isinstance(drm, Widevine) and not isinstance(self.cdm, WidevineCdm): - self.cdm = self.get_cdm(self.service, self.profile, drm="widevine") - elif isinstance(drm, PlayReady) and not isinstance(self.cdm, PlayReadyCdm): - self.cdm = self.get_cdm(self.service, self.profile, drm="playready") + if isinstance(drm, Widevine): + if not isinstance(self.cdm, (WidevineCdm, DecryptLabsRemoteCDM)) or ( + isinstance(self.cdm, DecryptLabsRemoteCDM) and self.cdm.is_playready + ): + widevine_cdm = self.get_cdm(self.service, self.profile, drm="widevine") + if widevine_cdm: + self.log.info("Switching to Widevine CDM for Widevine content") + self.cdm = widevine_cdm + elif isinstance(drm, PlayReady): + if not isinstance(self.cdm, (PlayReadyCdm, DecryptLabsRemoteCDM)) or ( + isinstance(self.cdm, DecryptLabsRemoteCDM) and not self.cdm.is_playready + ): + playready_cdm = self.get_cdm(self.service, self.profile, drm="playready") + if playready_cdm: + self.log.info("Switching to PlayReady CDM for PlayReady content") + self.cdm = playready_cdm if isinstance(drm, Widevine): with self.DRM_TABLE_LOCK: @@ -1486,8 +1506,8 @@ class dl: return Credential(*credentials) return Credential.loads(credentials) # type: ignore - @staticmethod def get_cdm( + self, service: str, profile: Optional[str] = None, drm: Optional[str] = None, @@ -1521,10 +1541,18 @@ class dl: cdm_api = next(iter(x for x in config.remote_cdm if x["name"] == cdm_name), None) if cdm_api: - is_decrypt_lab = True if cdm_api["type"] == "decrypt_labs" else False - del cdm_api["name"] - del cdm_api["type"] - return DecryptLabsRemoteCDM(service_name=service, **cdm_api) if is_decrypt_lab else RemoteCdm(**cdm_api) + is_decrypt_lab = True if cdm_api.get("type") == "decrypt_labs" else False + if is_decrypt_lab: + del cdm_api["name"] + del cdm_api["type"] + + # All DecryptLabs CDMs use DecryptLabsRemoteCDM + return DecryptLabsRemoteCDM(service_name=service, vaults=self.vaults, **cdm_api) + else: + del cdm_api["name"] + if "type" in cdm_api: + del cdm_api["type"] + return RemoteCdm(**cdm_api) prd_path = config.directories.prds / f"{cdm_name}.prd" if not prd_path.is_file(): @@ -1547,4 +1575,3 @@ class dl: raise ValueError(f"{cdm_name}.wvd is an invalid or corrupt Widevine Device file, {e}") return WidevineCdm.from_device(device) - diff --git a/packages/envied/src/envied/core/__init__.py b/packages/envied/src/envied/core/__init__.py index daa50c7..c0f285b 100644 --- a/packages/envied/src/envied/core/__init__.py +++ b/packages/envied/src/envied/core/__init__.py @@ -1 +1 @@ -__version__ = "1.4.2" +__version__ = "1.4.4" diff --git a/packages/envied/src/envied/core/__main__.py b/packages/envied/src/envied/core/__main__.py index 2c0b49b..beccf1d 100644 --- a/packages/envied/src/envied/core/__main__.py +++ b/packages/envied/src/envied/core/__main__.py @@ -10,13 +10,13 @@ from rich.padding import Padding from rich.text import Text from urllib3.exceptions import InsecureRequestWarning -from . import __version__ -from .commands import Commands -from .config import config -from .console import ComfyRichHandler, console -from .constants import context_settings -from .update_checker import UpdateChecker -from .utilities import rotate_log_file +from envied.core import __version__ +from envied.core.commands import Commands +from envied.core.config import config +from envied.core.console import ComfyRichHandler, console +from envied.core.constants import context_settings +from envied.core.update_checker import UpdateChecker +from envied.core.utilities import rotate_log_file LOGGING_PATH = None @@ -67,15 +67,35 @@ def main(version: bool, debug: bool, log_path: Path) -> None: r"░▀▀▀░▀░▀░░▀░░▀▀▀░▀▀▀░▀▀░" + "\n" , style="ascii.art", ), - "v 3.3.3 Copyright © 2019-2025 rlaphoenix" - + f"\nv [repr.number]{__version__}[/] - unshackle", + + f"\nv [repr.number]{__version__}[/] - https://github.com/unshackle-dl/unshackle", ), (1, 11, 1, 10), expand=True, ), justify="center", ) - + + + '''if version: + return + + if config.update_checks: + try: + latest_version = UpdateChecker.check_for_updates_sync(__version__) + if latest_version: + console.print( + f"\n[yellow]⚠️ Update available![/yellow] " + f"Current: {__version__} → Latest: [green]{latest_version}[/green]", + justify="center", + ) + console.print( + "Visit: https://github.com/envied-dl/envied/releases/latest\n", + justify="center", + ) + except Exception: + pass + ''' @atexit.register def save_log(): diff --git a/packages/envied/src/envied/core/cdm/decrypt_labs_remote_cdm.py b/packages/envied/src/envied/core/cdm/decrypt_labs_remote_cdm.py index c040362..0e005b8 100644 --- a/packages/envied/src/envied/core/cdm/decrypt_labs_remote_cdm.py +++ b/packages/envied/src/envied/core/cdm/decrypt_labs_remote_cdm.py @@ -1,143 +1,658 @@ +from __future__ import annotations + import base64 import secrets -from typing import Optional, Type, Union +from typing import Any, Dict, List, Optional, Union from uuid import UUID import requests -from pywidevine import PSSH, Device, DeviceTypes, Key, RemoteCdm -from pywidevine.license_protocol_pb2 import SignedDrmCertificate, SignedMessage +from pywidevine.device import DeviceTypes +from requests import Session -# Copyright 2024 by DevYukine. +from envied.core.vaults import Vaults -class DecryptLabsRemoteCDM(RemoteCdm): +class MockCertificateChain: + """Mock certificate chain for PlayReady compatibility.""" + + def __init__(self, name: str): + self._name = name + + def get_name(self) -> str: + return self._name + + +class Key: + """Key object compatible with pywidevine.""" + + def __init__(self, kid: str, key: str, type_: str = "CONTENT"): + if isinstance(kid, str): + clean_kid = kid.replace("-", "") + if len(clean_kid) == 32: + self.kid = UUID(hex=clean_kid) + else: + self.kid = UUID(hex=clean_kid.ljust(32, "0")) + else: + self.kid = kid + + if isinstance(key, str): + self.key = bytes.fromhex(key) + else: + self.key = key + + self.type = type_ + + +class DecryptLabsRemoteCDMExceptions: + """Exception classes for compatibility with pywidevine CDM.""" + + class InvalidSession(Exception): + """Raised when session ID is invalid.""" + + class TooManySessions(Exception): + """Raised when session limit is reached.""" + + class InvalidInitData(Exception): + """Raised when PSSH/init data is invalid.""" + + class InvalidLicenseType(Exception): + """Raised when license type is invalid.""" + + class InvalidLicenseMessage(Exception): + """Raised when license message is invalid.""" + + class InvalidContext(Exception): + """Raised when session has no context data.""" + + class SignatureMismatch(Exception): + """Raised when signature verification fails.""" + + +class DecryptLabsRemoteCDM: + """ + Decrypt Labs Remote CDM implementation with intelligent caching system. + + This class provides a drop-in replacement for pywidevine's local CDM using + Decrypt Labs' KeyXtractor API service, enhanced with smart caching logic + that minimizes unnecessary license requests. + + Key Features: + - Compatible with both Widevine and PlayReady DRM schemes + - Intelligent caching that compares required vs. available keys + - Automatic key combination for mixed cache/license scenarios + - Seamless fallback to license requests when keys are missing + + Intelligent Caching System: + 1. DRM classes (PlayReady/Widevine) provide required KIDs via set_required_kids() + 2. get_license_challenge() first checks for cached keys + 3. If cached keys satisfy requirements, returns empty challenge (no license needed) + 4. If keys are missing, makes targeted license request for remaining keys + 5. parse_license() combines cached and license keys intelligently + """ + + service_certificate_challenge = b"\x08\x04" + def __init__( self, - device_type: Union[DeviceTypes, str], - system_id: int, - security_level: int, - host: str, secret: str, - device_name: str, - service_name: str, + host: str = "https://keyxtractor.decryptlabs.com", + device_name: str = "ChromeCDM", + service_name: Optional[str] = None, + vaults: Optional[Vaults] = None, + device_type: Optional[str] = None, + system_id: Optional[int] = None, + security_level: Optional[int] = None, + **kwargs, ): - self.response_counter = 0 - self.pssh = None - self.api_session_ids = {} - self.license_request = None - self.service_name = service_name - self.keys = {} - try: - super().__init__(device_type, system_id, security_level, host, secret, device_name) - except Exception: - pass - self.req_session = requests.Session() - self.req_session.headers.update({"decrypt-labs-api-key": secret}) + """ + Initialize Decrypt Labs Remote CDM for Widevine and PlayReady schemes. - @classmethod - def from_device(cls, device: Device) -> Type["DecryptLabsRemoteCDM"]: - raise NotImplementedError("You cannot load a DecryptLabsRemoteCDM from a local Device file.") + Args: + secret: Decrypt Labs API key (matches config format) + host: Decrypt Labs API host URL (matches config format) + device_name: DRM scheme (ChromeCDM, L1, L2 for Widevine; SL2, SL3 for PlayReady) + service_name: Service name for key caching and vault operations + vaults: Vaults instance for local key caching + device_type: Device type (CHROME, ANDROID, PLAYREADY) - for compatibility + system_id: System ID - for compatibility + security_level: Security level - for compatibility + """ + _ = kwargs - def open(self) -> bytes: - # We stub this method to return a random session ID for now, later we save the api session id and resolve by our random generated one. - return bytes.fromhex(secrets.token_hex(16)) + self.secret = secret + self.host = host.rstrip("/") + self.device_name = device_name + self.service_name = service_name or "" + self.vaults = vaults + self.uch = self.host != "https://keyxtractor.decryptlabs.com" - def close(self, session_id: bytes) -> None: - # We stub this method to do nothing. - pass + self._device_type_str = device_type + if device_type: + self.device_type = self._get_device_type_enum(device_type) - def set_service_certificate(self, session_id: bytes, certificate: Optional[Union[bytes, str]]) -> str: - if isinstance(certificate, bytes): - certificate = base64.b64encode(certificate).decode() + self._is_playready = (device_type and device_type.upper() == "PLAYREADY") or (device_name in ["SL2", "SL3"]) - # certificate needs to be base64 to be sent off to the API. - # it needs to intentionally be kept as base64 encoded SignedMessage. + if self._is_playready: + self.system_id = system_id or 0 + self.security_level = security_level or (2000 if device_name == "SL2" else 3000) + else: + self.system_id = system_id or 26830 + self.security_level = security_level or 3 - self.req_session.signed_device_certificate = certificate - self.req_session.privacy_mode = True - - return "success" - - def get_service_certificate(self, session_id: bytes) -> Optional[SignedDrmCertificate]: - raise NotImplementedError("This method is not implemented in this CDM") - - def get_license_challenge( - self, session_id: bytes, pssh: PSSH, license_type: str = "STREAMING", privacy_mode: bool = True - ) -> bytes: - self.pssh = pssh - - res = self.session( - self.host + "/get-request", + self._sessions: Dict[bytes, Dict[str, Any]] = {} + self._pssh_b64 = None + self._required_kids: Optional[List[str]] = None + self._http_session = Session() + self._http_session.headers.update( { - "init_data": self.pssh.dumps(), - "service_certificate": self.req_session.signed_device_certificate, - "scheme": "widevine", - "service": self.service_name, - }, + "decrypt-labs-api-key": self.secret, + "Content-Type": "application/json", + "User-Agent": "envied-decrypt-labs-cdm/1.0", + } ) - self.license_request = res["challenge"] - self.api_session_ids[session_id] = res["session_id"] - - return base64.b64decode(self.license_request) - - def parse_license(self, session_id: bytes, license_message: Union[SignedMessage, bytes, str]) -> None: - session_id_api = self.api_session_ids[session_id] - if session_id not in self.keys: - self.keys[session_id] = [] - session_keys = self.keys[session_id] - - if isinstance(license_message, dict) and "keys" in license_message: - session_keys.extend( - [ - Key(kid=Key.kid_to_uuid(x["kid"]), type_=x.get("type", "CONTENT"), key=bytes.fromhex(x["key"])) - for x in license_message["keys"] - ] - ) - + def _get_device_type_enum(self, device_type: str): + """Convert device type string to enum for compatibility.""" + device_type_upper = device_type.upper() + if device_type_upper == "ANDROID": + return DeviceTypes.ANDROID + elif device_type_upper == "CHROME": + return DeviceTypes.CHROME else: - res = self.session( - self.host + "/decrypt-response", - { - "session_id": session_id_api, - "init_data": self.pssh.dumps(), - "license_request": self.license_request, - "license_response": license_message, - "scheme": "widevine", - }, - ) + return DeviceTypes.CHROME - original_keys = res["keys"].replace("\n", " ") - keys_separated = original_keys.split("--key ") - formatted_keys = [] - for k in keys_separated: - if ":" in k: - key = k.strip() - formatted_keys.append(key) - for keys in formatted_keys: - session_keys.append( - ( - Key( - kid=UUID(bytes=bytes.fromhex(keys.split(":")[0])), - type_="CONTENT", - key=bytes.fromhex(keys.split(":")[1]), - ) - ) + @property + def is_playready(self) -> bool: + """Check if this CDM is in PlayReady mode.""" + return self._is_playready + + @property + def certificate_chain(self) -> MockCertificateChain: + """Mock certificate chain for PlayReady compatibility.""" + return MockCertificateChain(f"{self.device_name}_Remote") + + def set_pssh_b64(self, pssh_b64: str) -> None: + """Store base64-encoded PSSH data for PlayReady compatibility.""" + self._pssh_b64 = pssh_b64 + + def set_required_kids(self, kids: List[Union[str, UUID]]) -> None: + """ + Set the required Key IDs for intelligent caching decisions. + + This method enables the CDM to make smart decisions about when to request + additional keys via license challenges. When cached keys are available, + the CDM will compare them against the required KIDs to determine if a + license request is still needed for missing keys. + + Args: + kids: List of required Key IDs as UUIDs or hex strings + + Note: + Should be called by DRM classes (PlayReady/Widevine) before making + license challenge requests to enable optimal caching behavior. + """ + self._required_kids = [] + for kid in kids: + if isinstance(kid, UUID): + self._required_kids.append(str(kid).replace("-", "").lower()) + else: + self._required_kids.append(str(kid).replace("-", "").lower()) + + def _generate_session_id(self) -> bytes: + """Generate a unique session ID.""" + return secrets.token_bytes(16) + + def _get_init_data_from_pssh(self, pssh: Any) -> str: + """Extract init data from various PSSH formats.""" + if self.is_playready and self._pssh_b64: + return self._pssh_b64 + + if hasattr(pssh, "dumps"): + dumps_result = pssh.dumps() + + if isinstance(dumps_result, str): + try: + base64.b64decode(dumps_result) + return dumps_result + except Exception: + return base64.b64encode(dumps_result.encode("utf-8")).decode("utf-8") + else: + return base64.b64encode(dumps_result).decode("utf-8") + elif hasattr(pssh, "raw"): + raw_data = pssh.raw + if isinstance(raw_data, str): + raw_data = raw_data.encode("utf-8") + return base64.b64encode(raw_data).decode("utf-8") + elif hasattr(pssh, "__class__") and "WrmHeader" in pssh.__class__.__name__: + if self.is_playready: + raise ValueError("PlayReady WRM header received but no PSSH B64 was set via set_pssh_b64()") + + if hasattr(pssh, "raw_bytes"): + return base64.b64encode(pssh.raw_bytes).decode("utf-8") + elif hasattr(pssh, "bytes"): + return base64.b64encode(pssh.bytes).decode("utf-8") + else: + raise ValueError(f"Cannot extract PSSH data from WRM header type: {type(pssh)}") + else: + raise ValueError(f"Unsupported PSSH type: {type(pssh)}") + + def open(self) -> bytes: + """ + Open a new CDM session. + + Returns: + Session identifier as bytes + """ + session_id = self._generate_session_id() + self._sessions[session_id] = { + "service_certificate": None, + "keys": [], + "pssh": None, + "challenge": None, + "decrypt_labs_session_id": None, + } + return session_id + + def close(self, session_id: bytes) -> None: + """ + Close a CDM session. + + Args: + session_id: Session identifier + + Raises: + ValueError: If session ID is invalid + """ + if session_id not in self._sessions: + raise DecryptLabsRemoteCDMExceptions.InvalidSession(f"Invalid session ID: {session_id.hex()}") + + del self._sessions[session_id] + + def get_service_certificate(self, session_id: bytes) -> Optional[bytes]: + """ + Get the service certificate for a session. + + Args: + session_id: Session identifier + + Returns: + Service certificate if set, None otherwise + + Raises: + ValueError: If session ID is invalid + """ + if session_id not in self._sessions: + raise DecryptLabsRemoteCDMExceptions.InvalidSession(f"Invalid session ID: {session_id.hex()}") + + return self._sessions[session_id]["service_certificate"] + + def set_service_certificate(self, session_id: bytes, certificate: Optional[Union[bytes, str]]) -> str: + """ + Set the service certificate for a session. + + Args: + session_id: Session identifier + certificate: Service certificate (bytes or base64 string) + + Returns: + Certificate status message + + Raises: + ValueError: If session ID is invalid + """ + if session_id not in self._sessions: + raise DecryptLabsRemoteCDMExceptions.InvalidSession(f"Invalid session ID: {session_id.hex()}") + + if certificate is None: + self._sessions[session_id]["service_certificate"] = None + return "Removed" + + if isinstance(certificate, str): + certificate = base64.b64decode(certificate) + + self._sessions[session_id]["service_certificate"] = certificate + return "Successfully set Service Certificate" + + def has_cached_keys(self, session_id: bytes) -> bool: + """ + Check if cached keys are available for the session. + + Args: + session_id: Session identifier + + Returns: + True if cached keys are available + + Raises: + ValueError: If session ID is invalid + """ + if session_id not in self._sessions: + raise DecryptLabsRemoteCDMExceptions.InvalidSession(f"Invalid session ID: {session_id.hex()}") + + session = self._sessions[session_id] + session_keys = session.get("keys", []) + return len(session_keys) > 0 + + def get_license_challenge( + self, session_id: bytes, pssh_or_wrm: Any, license_type: str = "STREAMING", privacy_mode: bool = True + ) -> bytes: + """ + Generate a license challenge using Decrypt Labs API with intelligent caching. + + This method implements smart caching logic that: + 1. First attempts to retrieve cached keys from the API + 2. If required KIDs are set, compares cached keys against requirements + 3. Only makes a license request if keys are missing + 4. Returns empty challenge if all required keys are cached + + The intelligent caching works as follows: + - With required KIDs set: Only requests license for missing keys + - Without required KIDs: Returns any available cached keys + - For PlayReady: Combines cached keys with license keys seamlessly + + Args: + session_id: Session identifier + pssh_or_wrm: PSSH object or WRM header (for PlayReady compatibility) + license_type: Type of license (STREAMING, OFFLINE, AUTOMATIC) - for compatibility only + privacy_mode: Whether to use privacy mode - for compatibility only + + Returns: + License challenge as bytes, or empty bytes if cached keys satisfy requirements + + Raises: + InvalidSession: If session ID is invalid + requests.RequestException: If API request fails + + Note: + Call set_required_kids() before this method for optimal caching behavior. + """ + _ = license_type, privacy_mode + + if session_id not in self._sessions: + raise DecryptLabsRemoteCDMExceptions.InvalidSession(f"Invalid session ID: {session_id.hex()}") + + session = self._sessions[session_id] + + session["pssh"] = pssh_or_wrm + init_data = self._get_init_data_from_pssh(pssh_or_wrm) + already_tried_cache = session.get("tried_cache", False) + + request_data = { + "scheme": self.device_name, + "init_data": init_data, + "get_cached_keys_if_exists": not already_tried_cache, + } + + if self.device_name in ["L1", "L2", "SL2", "SL3"] and self.service_name: + request_data["service"] = self.service_name + + if session["service_certificate"]: + request_data["service_certificate"] = base64.b64encode(session["service_certificate"]).decode("utf-8") + + response = self._http_session.post(f"{self.host}/get-request", json=request_data, timeout=30) + + if response.status_code != 200: + raise requests.RequestException(f"API request failed: {response.status_code} {response.text}") + + data = response.json() + + if data.get("message") != "success": + error_msg = data.get("message", "Unknown error") + if "details" in data: + error_msg += f" - Details: {data['details']}" + if "error" in data: + error_msg += f" - Error: {data['error']}" + + if "service_certificate is required" in str(data) and not session["service_certificate"]: + error_msg += " (No service certificate was provided to the CDM session)" + + raise requests.RequestException(f"API error: {error_msg}") + + message_type = data.get("message_type") + + if message_type == "cached-keys" or "cached_keys" in data: + """ + Handle cached keys response from API. + + When the API returns cached keys, we need to determine if they satisfy + our requirements or if we need to make an additional license request + for missing keys. + """ + cached_keys = data.get("cached_keys", []) + parsed_keys = self._parse_cached_keys(cached_keys) + session["keys"] = parsed_keys + session["tried_cache"] = True + + if self._required_kids: + cached_kids = set() + for key in parsed_keys: + if isinstance(key, dict) and "kid" in key: + cached_kids.add(key["kid"].replace("-", "").lower()) + + required_kids = set(self._required_kids) + missing_kids = required_kids - cached_kids + + if missing_kids: + session["cached_keys"] = parsed_keys + request_data["get_cached_keys_if_exists"] = False + response = self._http_session.post(f"{self.host}/get-request", json=request_data, timeout=30) + if response.status_code == 200: + data = response.json() + if data.get("message") == "success" and "challenge" in data: + challenge = base64.b64decode(data["challenge"]) + session["challenge"] = challenge + session["decrypt_labs_session_id"] = data["session_id"] + return challenge + + return b"" + else: + return b"" + else: + return b"" + + if message_type == "license-request" or "challenge" in data: + challenge = base64.b64decode(data["challenge"]) + session["challenge"] = challenge + session["decrypt_labs_session_id"] = data["session_id"] + return challenge + + error_msg = f"Unexpected API response format. message_type={message_type}, available_fields={list(data.keys())}" + if data.get("message"): + error_msg = f"API response: {data['message']} - {error_msg}" + if "details" in data: + error_msg += f" - Details: {data['details']}" + if "error" in data: + error_msg += f" - Error: {data['error']}" + + if already_tried_cache and data.get("message") == "success": + return b"" + + raise requests.RequestException(error_msg) + + def parse_license(self, session_id: bytes, license_message: Union[bytes, str]) -> None: + """ + Parse license response using Decrypt Labs API with intelligent key combination. + + For PlayReady content with partial cached keys, this method intelligently + combines the cached keys with newly obtained license keys, avoiding + duplicates while ensuring all required keys are available. + + The key combination process: + 1. Extracts keys from the license response + 2. If cached keys exist (PlayReady), combines them with license keys + 3. Removes duplicate keys by comparing normalized KIDs + 4. Updates the session with the complete key set + + Args: + session_id: Session identifier + license_message: License response from license server + + Raises: + ValueError: If session ID is invalid or no challenge available + requests.RequestException: If API request fails + """ + if session_id not in self._sessions: + raise DecryptLabsRemoteCDMExceptions.InvalidSession(f"Invalid session ID: {session_id.hex()}") + + session = self._sessions[session_id] + + if session["keys"] and not (self.is_playready and "cached_keys" in session): + return + + if not session.get("challenge") or not session.get("decrypt_labs_session_id"): + raise ValueError("No challenge available - call get_license_challenge first") + + if isinstance(license_message, str): + if self.is_playready and license_message.strip().startswith(" List[Key]: + """ + Get keys from the session. + + Args: + session_id: Session identifier + type_: Optional key type filter (CONTENT, SIGNING, etc.) + + Returns: + List of Key objects + + Raises: + InvalidSession: If session ID is invalid + """ + if session_id not in self._sessions: + raise DecryptLabsRemoteCDMExceptions.InvalidSession(f"Invalid session ID: {session_id.hex()}") + + key_dicts = self._sessions[session_id]["keys"] + keys = [Key(kid=k["kid"], key=k["key"], type_=k["type"]) for k in key_dicts] + + if type_: + keys = [key for key in keys if key.type == type_] + + return keys + + def _parse_cached_keys(self, cached_keys_data: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """Parse cached keys from API response. + + Args: + cached_keys_data: List of cached key objects from API + + Returns: + List of key dictionaries + """ + keys = [] + + try: + if cached_keys_data and isinstance(cached_keys_data, list): + for key_data in cached_keys_data: + if "kid" in key_data and "key" in key_data: + keys.append({"kid": key_data["kid"], "key": key_data["key"], "type": "CONTENT"}) + except Exception: + pass + return keys + + def _parse_keys_response(self, data: Dict[str, Any]) -> List[Dict[str, Any]]: + """Parse keys from decrypt response.""" + keys = [] + + if "keys" in data and isinstance(data["keys"], str): + keys_string = data["keys"] + + for line in keys_string.split("\n"): + line = line.strip() + if line.startswith("--key "): + key_part = line[6:] + if ":" in key_part: + kid, key = key_part.split(":", 1) + keys.append({"kid": kid.strip(), "key": key.strip(), "type": "CONTENT"}) + elif "keys" in data and isinstance(data["keys"], list): + for key_data in data["keys"]: + keys.append( + {"kid": key_data.get("kid"), "key": key_data.get("key"), "type": key_data.get("type", "CONTENT")} ) - def get_keys(self, session_id: bytes, type_: Optional[Union[int, str]] = None) -> list[Key]: - return self.keys[session_id] + return keys - def session(self, url, data, retries=3): - res = self.req_session.post(url, json=data).json() - if res.get("message") != "success": - if "License Response Decryption Process Failed at the very beginning" in res.get("Error", ""): - if retries > 0: - return self.session(url, data, retries=retries - 1) - else: - raise ValueError(f"CDM API returned an error: {res['Error']}") - else: - raise ValueError(f"CDM API returned an error: {res['Error']}") - - return res +__all__ = ["DecryptLabsRemoteCDM"] diff --git a/packages/envied/src/envied/core/commands.py b/packages/envied/src/envied/core/commands.py index e6024df..7f0fb52 100644 --- a/packages/envied/src/envied/core/commands.py +++ b/packages/envied/src/envied/core/commands.py @@ -2,8 +2,8 @@ from typing import Optional import click -from .config import config -from .utilities import import_module_by_path +from envied.core.config import config +from envied.core.utilities import import_module_by_path _COMMANDS = sorted( (path for path in config.directories.commands.glob("*.py") if path.stem.lower() != "__init__"), key=lambda x: x.stem diff --git a/packages/envied/src/envied/core/config.py b/packages/envied/src/envied/core/config.py index af0e79d..2967479 100644 --- a/packages/envied/src/envied/core/config.py +++ b/packages/envied/src/envied/core/config.py @@ -92,8 +92,7 @@ class Config: self.update_check_interval: int = kwargs.get("update_check_interval", 24) self.scene_naming: bool = kwargs.get("scene_naming", True) self.series_year: bool = kwargs.get("series_year", True) - self.pssh_display: str = kwargs.get("pssh_display") or "fold" - + self.pssh_display: str = kwargs.get("pssh_display") or "fold" self.title_cache_time: int = kwargs.get("title_cache_time", 1800) # 30 minutes default self.title_cache_max_retention: int = kwargs.get("title_cache_max_retention", 86400) # 24 hours default self.title_cache_enabled: bool = kwargs.get("title_cache_enabled", True) diff --git a/packages/envied/src/envied/core/drm/playready.py b/packages/envied/src/envied/core/drm/playready.py index 2734fc8..ef9c90e 100644 --- a/packages/envied/src/envied/core/drm/playready.py +++ b/packages/envied/src/envied/core/drm/playready.py @@ -224,41 +224,79 @@ class PlayReady: def kids(self) -> list[UUID]: return self._kids - def get_content_keys(self, cdm: PlayReadyCdm, certificate: Callable, licence: Callable) -> None: - for kid in self.kids: - if kid in self.content_keys: + def _extract_keys_from_cdm(self, cdm: PlayReadyCdm, session_id: bytes) -> dict: + """Extract keys from CDM session with cross-library compatibility. + + Args: + cdm: CDM instance + session_id: Session identifier + + Returns: + Dictionary mapping KID UUIDs to hex keys + """ + keys = {} + for key in cdm.get_keys(session_id): + if hasattr(key, "key_id"): + kid = key.key_id + elif hasattr(key, "kid"): + kid = key.kid + else: continue - session_id = cdm.open() - try: - challenge = cdm.get_license_challenge(session_id, self.pssh.wrm_headers[0]) - license_res = licence(challenge=challenge) - if isinstance(license_res, bytes): - license_str = license_res.decode(errors="ignore") - else: - license_str = str(license_res) + if hasattr(key, "key") and hasattr(key.key, "hex"): + key_hex = key.key.hex() + elif hasattr(key, "key") and isinstance(key.key, bytes): + key_hex = key.key.hex() + elif hasattr(key, "key") and isinstance(key.key, str): + key_hex = key.key + else: + continue - if "" not in license_str: - try: - license_str = base64.b64decode(license_str + "===").decode() - except Exception: - pass + keys[kid] = key_hex + return keys - cdm.parse_license(session_id, license_str) - keys = {key.key_id: key.key.hex() for key in cdm.get_keys(session_id)} - self.content_keys.update(keys) - finally: - cdm.close(session_id) + def get_content_keys(self, cdm: PlayReadyCdm, certificate: Callable, licence: Callable) -> None: + session_id = cdm.open() + try: + if hasattr(cdm, "set_pssh_b64") and self.pssh_b64: + cdm.set_pssh_b64(self.pssh_b64) + + if hasattr(cdm, "set_required_kids"): + cdm.set_required_kids(self.kids) + + challenge = cdm.get_license_challenge(session_id, self.pssh.wrm_headers[0]) + + if challenge: + try: + license_res = licence(challenge=challenge) + if isinstance(license_res, bytes): + license_str = license_res.decode(errors="ignore") + else: + license_str = str(license_res) + + if "" not in license_str: + try: + license_str = base64.b64decode(license_str + "===").decode() + except Exception: + pass + + cdm.parse_license(session_id, license_str) + except Exception: + raise + + keys = self._extract_keys_from_cdm(cdm, session_id) + self.content_keys.update(keys) + finally: + cdm.close(session_id) if not self.content_keys: raise PlayReady.Exceptions.EmptyLicense("No Content Keys were within the License") - def decrypt(self, path: Path, use_mp4decrypt: bool = False) -> None: + def decrypt(self, path: Path) -> None: """ Decrypt a Track with PlayReady DRM. Args: path: Path to the encrypted file to decrypt - use_mp4decrypt: If True, use mp4decrypt instead of Shaka Packager Raises: EnvironmentError if the required decryption executable could not be found. ValueError if the track has not yet been downloaded. @@ -270,7 +308,9 @@ class PlayReady: if not path or not path.exists(): raise ValueError("Tried to decrypt a file that does not exist.") - if use_mp4decrypt: + decrypter = str(getattr(config, "decryption", "")).lower() + + if decrypter == "mp4decrypt": return self._decrypt_with_mp4decrypt(path) else: return self._decrypt_with_shaka_packager(path) diff --git a/packages/envied/src/envied/core/drm/widevine.py b/packages/envied/src/envied/core/drm/widevine.py index 3b1bdb6..f8ff998 100644 --- a/packages/envied/src/envied/core/drm/widevine.py +++ b/packages/envied/src/envied/core/drm/widevine.py @@ -185,7 +185,15 @@ class Widevine: if cert and hasattr(cdm, "set_service_certificate"): cdm.set_service_certificate(session_id, cert) - cdm.parse_license(session_id, licence(challenge=cdm.get_license_challenge(session_id, self.pssh))) + if hasattr(cdm, "set_required_kids"): + cdm.set_required_kids(self.kids) + + challenge = cdm.get_license_challenge(session_id, self.pssh) + + if hasattr(cdm, "has_cached_keys") and cdm.has_cached_keys(session_id): + pass + else: + cdm.parse_license(session_id, licence(challenge=challenge)) self.content_keys = {key.kid: key.key.hex() for key in cdm.get_keys(session_id, "CONTENT")} if not self.content_keys: @@ -213,10 +221,18 @@ class Widevine: if cert and hasattr(cdm, "set_service_certificate"): cdm.set_service_certificate(session_id, cert) - cdm.parse_license( - session_id, - licence(session_id=session_id, challenge=cdm.get_license_challenge(session_id, self.pssh)), - ) + if hasattr(cdm, "set_required_kids"): + cdm.set_required_kids(self.kids) + + challenge = cdm.get_license_challenge(session_id, self.pssh) + + if hasattr(cdm, "has_cached_keys") and cdm.has_cached_keys(session_id): + pass + else: + cdm.parse_license( + session_id, + licence(session_id=session_id, challenge=challenge), + ) self.content_keys = {key.kid: key.key.hex() for key in cdm.get_keys(session_id, "CONTENT")} if not self.content_keys: @@ -227,12 +243,11 @@ class Widevine: finally: cdm.close(session_id) - def decrypt(self, path: Path, use_mp4decrypt: bool = False) -> None: + def decrypt(self, path: Path) -> None: """ Decrypt a Track with Widevine DRM. Args: path: Path to the encrypted file to decrypt - use_mp4decrypt: If True, use mp4decrypt instead of Shaka Packager Raises: EnvironmentError if the required decryption executable could not be found. ValueError if the track has not yet been downloaded. @@ -244,7 +259,9 @@ class Widevine: if not path or not path.exists(): raise ValueError("Tried to decrypt a file that does not exist.") - if use_mp4decrypt: + decrypter = str(getattr(config, "decryption", "")).lower() + + if decrypter == "mp4decrypt": return self._decrypt_with_mp4decrypt(path) else: return self._decrypt_with_shaka_packager(path) diff --git a/packages/envied/src/envied/core/tracks/track.py b/packages/envied/src/envied/core/tracks/track.py index c12323e..dbdd1ca 100644 --- a/packages/envied/src/envied/core/tracks/track.py +++ b/packages/envied/src/envied/core/tracks/track.py @@ -420,6 +420,15 @@ class Track: for drm in self.drm: if isinstance(drm, PlayReady): return drm + elif hasattr(cdm, 'is_playready'): + if cdm.is_playready: + for drm in self.drm: + if isinstance(drm, PlayReady): + return drm + else: + for drm in self.drm: + if isinstance(drm, Widevine): + return drm return self.drm[0] diff --git a/packages/envied/src/envied/core/utilities.py b/packages/envied/src/envied/core/utilities.py index f0b1cb0..22ab989 100644 --- a/packages/envied/src/envied/core/utilities.py +++ b/packages/envied/src/envied/core/utilities.py @@ -22,9 +22,9 @@ from langcodes import Language, closest_match from pymp4.parser import Box from unidecode import unidecode -from .cacher import Cacher -from .config import config -from .constants import LANGUAGE_MAX_DISTANCE +from envied.core.cacher import Cacher +from envied.core.config import config +from envied.core.constants import LANGUAGE_MAX_DISTANCE def rotate_log_file(log_path: Path, keep: int = 20) -> Path: diff --git a/packages/envied/src/envied/core/utils/tags.py b/packages/envied/src/envied/core/utils/tags.py index 2ba2705..f607009 100644 --- a/packages/envied/src/envied/core/utils/tags.py +++ b/packages/envied/src/envied/core/utils/tags.py @@ -8,6 +8,7 @@ import tempfile from difflib import SequenceMatcher from pathlib import Path from typing import Optional, Tuple +from xml.sax.saxutils import escape import requests from requests.adapters import HTTPAdapter, Retry @@ -289,9 +290,9 @@ def _apply_tags(path: Path, tags: dict[str, str]) -> None: log.debug("mkvpropedit not found on PATH; skipping tags") return log.debug("Applying tags to %s: %s", path, tags) - xml_lines = ["", "", " ", " "] + xml_lines = ['', "", " ", " "] for name, value in tags.items(): - xml_lines.append(f" {name}{value}") + xml_lines.append(f" {escape(name)}{escape(value)}") xml_lines.extend([" ", ""]) with tempfile.NamedTemporaryFile("w", suffix=".xml", delete=False) as f: f.write("\n".join(xml_lines)) @@ -349,13 +350,25 @@ def tag_file(path: Path, title: Title, tmdb_id: Optional[int] | None = None) -> if simkl_tmdb_id: tmdb_id = simkl_tmdb_id - show_ids = simkl_data.get("show", {}).get("ids", {}) - if show_ids.get("imdb"): - standard_tags["IMDB"] = f"https://www.imdb.com/title/{show_ids['imdb']}" - if show_ids.get("tvdb"): - standard_tags["TVDB"] = f"https://thetvdb.com/dereferrer/series/{show_ids['tvdb']}" - if show_ids.get("tmdbtv"): - standard_tags["TMDB"] = f"https://www.themoviedb.org/tv/{show_ids['tmdbtv']}" + # Handle TV show data from Simkl + if simkl_data.get("type") == "episode" and "show" in simkl_data: + show_ids = simkl_data.get("show", {}).get("ids", {}) + if show_ids.get("imdb"): + standard_tags["IMDB"] = show_ids["imdb"] + if show_ids.get("tvdb"): + standard_tags["TVDB2"] = f"series/{show_ids['tvdb']}" + if show_ids.get("tmdbtv"): + standard_tags["TMDB"] = f"tv/{show_ids['tmdbtv']}" + + # Handle movie data from Simkl + elif simkl_data.get("type") == "movie" and "movie" in simkl_data: + movie_ids = simkl_data.get("movie", {}).get("ids", {}) + if movie_ids.get("imdb"): + standard_tags["IMDB"] = movie_ids["imdb"] + if movie_ids.get("tvdb"): + standard_tags["TVDB2"] = f"movies/{movie_ids['tvdb']}" + if movie_ids.get("tmdb"): + standard_tags["TMDB"] = f"movie/{movie_ids['tmdb']}" # Use TMDB API for additional metadata (either from provided ID or Simkl lookup) api_key = _api_key() @@ -373,8 +386,8 @@ def tag_file(path: Path, title: Title, tmdb_id: Optional[int] | None = None) -> _apply_tags(path, custom_tags) return - tmdb_url = f"https://www.themoviedb.org/{'movie' if kind == 'movie' else 'tv'}/{tmdb_id}" - standard_tags["TMDB"] = tmdb_url + prefix = "movie" if kind == "movie" else "tv" + standard_tags["TMDB"] = f"{prefix}/{tmdb_id}" try: ids = external_ids(tmdb_id, kind) except requests.RequestException as exc: @@ -385,11 +398,13 @@ def tag_file(path: Path, title: Title, tmdb_id: Optional[int] | None = None) -> imdb_id = ids.get("imdb_id") if imdb_id: - standard_tags["IMDB"] = f"https://www.imdb.com/title/{imdb_id}" + standard_tags["IMDB"] = imdb_id tvdb_id = ids.get("tvdb_id") if tvdb_id: - tvdb_prefix = "movies" if kind == "movie" else "series" - standard_tags["TVDB"] = f"https://thetvdb.com/dereferrer/{tvdb_prefix}/{tvdb_id}" + if kind == "movie": + standard_tags["TVDB2"] = f"movies/{tvdb_id}" + else: + standard_tags["TVDB2"] = f"series/{tvdb_id}" merged_tags = { **custom_tags, diff --git a/packages/envied/src/envied/services/NF/MSL/__init__.py b/packages/envied/src/envied/services/NF/MSL/__init__.py index e0ced87..5e1662e 100644 --- a/packages/envied/src/envied/services/NF/MSL/__init__.py +++ b/packages/envied/src/envied/services/NF/MSL/__init__.py @@ -26,7 +26,7 @@ from .schemes import EntityAuthenticationSchemes # noqa: F401 from .schemes import KeyExchangeSchemes from .schemes.EntityAuthentication import EntityAuthentication from .schemes.KeyExchangeRequest import KeyExchangeRequest -# from vinetrimmer.utils.widevine.device import RemoteDevice +# from vinetrimmer.utils.wienvied.device import RemoteDevice class MSL: log = logging.getLogger("MSL") diff --git a/packages/envied/src/envied/services/RTE/__init__.py b/packages/envied/src/envied/services/RTE/__init__.py index e255f80..9e3b3a5 100644 --- a/packages/envied/src/envied/services/RTE/__init__.py +++ b/packages/envied/src/envied/services/RTE/__init__.py @@ -45,7 +45,7 @@ class RTE(Service): """ - # GEOFENCE = ("ie",) + #GEOFENCE = ("ie",) @staticmethod @click.command(name="RTE", short_help="https://www.rte.ie/player/", help=__doc__) diff --git a/packages/envied/src/envied/services/RTE/config.yaml b/packages/envied/src/envied/services/RTE/config.yaml index cf6eb37..8b82d52 100644 --- a/packages/envied/src/envied/services/RTE/config.yaml +++ b/packages/envied/src/envied/services/RTE/config.yaml @@ -4,4 +4,4 @@ headers: endpoints: base_url: https://www.rte.ie feed: https://feed.entertainment.tv.theplatform.eu - license: https://widevine.entitlement.eu.theplatform.com/wv/web/ModularDrm + license: https://wienvied.entitlement.eu.theplatform.com/wv/web/ModularDrm diff --git a/packages/envied/src/envied/services/STV/__init__.py b/packages/envied/src/envied/services/STV/__init__.py index 8cbd9b4..059bf0d 100644 --- a/packages/envied/src/envied/services/STV/__init__.py +++ b/packages/envied/src/envied/services/STV/__init__.py @@ -166,11 +166,11 @@ class STV(Service): source for source in data["sources"] if source.get("type") == "application/dash+xml" - and source.get("key_systems").get("com.widevine.alpha")), + and source.get("key_systems").get("com.wienvied.alpha")), None, ) - self.license = key_systems["key_systems"]["com.widevine.alpha"]["license_url"] if key_systems else None + self.license = key_systems["key_systems"]["com.wienvied.alpha"]["license_url"] if key_systems else None manifest = self.trim_duration(source_manifest) tracks = DASH.from_text(manifest, source_manifest).to_tracks(title.language) diff --git a/packages/envied/src/envied/services/TPTV/__init__.py b/packages/envied/src/envied/services/TPTV/__init__.py index 29e7fed..483d966 100644 --- a/packages/envied/src/envied/services/TPTV/__init__.py +++ b/packages/envied/src/envied/services/TPTV/__init__.py @@ -20,7 +20,7 @@ def get_widevine_license_url(manifest_str): try: data = json.loads(manifest_str) for source in data.get("sources", []): - widevine = source.get("key_systems", {}).get("com.widevine.alpha") + widevine = source.get("key_systems", {}).get("com.wienvied.alpha") if widevine and "license_url" in widevine: return widevine["license_url"] except json.JSONDecodeError: diff --git a/packages/envied/src/envied/services/TVNZ/__init__.py b/packages/envied/src/envied/services/TVNZ/__init__.py index 154e3c3..9c0a97a 100644 --- a/packages/envied/src/envied/services/TVNZ/__init__.py +++ b/packages/envied/src/envied/services/TVNZ/__init__.py @@ -172,14 +172,14 @@ class TVNZ(Service): ) self.license = next(( - x["key_systems"]["com.widevine.alpha"]["license_url"] + x["key_systems"]["com.wienvied.alpha"]["license_url"] for x in data["sources"] - if x.get("key_systems").get("com.widevine.alpha")), + if x.get("key_systems").get("com.wienvied.alpha")), None, ) source_manifest = next(( x["src"] for x in data["sources"] - if x.get("key_systems").get("com.widevine.alpha")), + if x.get("key_systems").get("com.wienvied.alpha")), None, ) diff --git a/packages/envied/src/envied/services/UKTV/__init__.py b/packages/envied/src/envied/services/UKTV/__init__.py index 18613fe..4dd5b19 100644 --- a/packages/envied/src/envied/services/UKTV/__init__.py +++ b/packages/envied/src/envied/services/UKTV/__init__.py @@ -120,14 +120,14 @@ class UKTV(Service): data = r.json() self.license = next(( - x["key_systems"]["com.widevine.alpha"]["license_url"] + x["key_systems"]["com.wienvied.alpha"]["license_url"] for x in data["sources"] - if x.get("key_systems").get("com.widevine.alpha")), + if x.get("key_systems").get("com.wienvied.alpha")), None, ) source_manifest = next(( x["src"] for x in data["sources"] - if x.get("key_systems").get("com.widevine.alpha")), + if x.get("key_systems").get("com.wienvied.alpha")), None, ) if not self.license or not source_manifest: diff --git a/packages/envied/src/envied/vaults/API.py b/packages/envied/src/envied/vaults/API.py new file mode 100644 index 0000000..2cc5b4a --- /dev/null +++ b/packages/envied/src/envied/vaults/API.py @@ -0,0 +1,184 @@ +from typing import Iterator, Optional, Union +from uuid import UUID + +from requests import Session + +from envied.core import __version__ +from envied.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): + super().__init__(name, no_push) + self.uri = uri.rstrip("/") + self.session = Session() + self.session.headers.update({"User-Agent": f"envied v{__version__}"}) + self.session.headers.update({"Authorization": f"Bearer {token}"}) + + def get_key(self, kid: Union[UUID, str], service: str) -> Optional[str]: + if isinstance(kid, UUID): + kid = kid.hex + + data = self.session.get( + url=f"{self.uri}/{service.lower()}/{kid}", headers={"Accept": "application/json"} + ).json() + + code = int(data.get("code", 0)) + message = data.get("message") + error = { + 0: None, + 1: Exceptions.AuthRejected, + 2: Exceptions.TooManyRequests, + 3: Exceptions.ServiceTagInvalid, + 4: Exceptions.KeyIdInvalid, + }.get(code, ValueError) + + if error: + raise error(f"{message} ({code})") + + content_key = data.get("content_key") + if not content_key: + return None + + if not isinstance(content_key, str): + raise ValueError(f"Expected {content_key} to be {str}, was {type(content_key)}") + + return content_key + + def get_keys(self, service: str) -> Iterator[tuple[str, str]]: + page = 1 + + while True: + data = self.session.get( + url=f"{self.uri}/{service.lower()}", + params={"page": page, "total": 10}, + headers={"Accept": "application/json"}, + ).json() + + code = int(data.get("code", 0)) + message = data.get("message") + error = { + 0: None, + 1: Exceptions.AuthRejected, + 2: Exceptions.TooManyRequests, + 3: Exceptions.PageInvalid, + 4: Exceptions.ServiceTagInvalid, + }.get(code, ValueError) + + if error: + raise error(f"{message} ({code})") + + content_keys = data.get("content_keys") + if content_keys: + if not isinstance(content_keys, dict): + raise ValueError(f"Expected {content_keys} to be {dict}, was {type(content_keys)}") + + for key_id, key in content_keys.items(): + yield key_id, key + + pages = int(data["pages"]) + if pages <= page: + break + + page += 1 + + def add_key(self, service: str, kid: Union[UUID, str], key: str) -> bool: + if isinstance(kid, UUID): + kid = kid.hex + + data = self.session.post( + url=f"{self.uri}/{service.lower()}/{kid}", json={"content_key": key}, headers={"Accept": "application/json"} + ).json() + + code = int(data.get("code", 0)) + message = data.get("message") + error = { + 0: None, + 1: Exceptions.AuthRejected, + 2: Exceptions.TooManyRequests, + 3: Exceptions.ServiceTagInvalid, + 4: Exceptions.KeyIdInvalid, + 5: Exceptions.ContentKeyInvalid, + }.get(code, ValueError) + + if error: + raise error(f"{message} ({code})") + + # the kid:key was new to the vault (optional) + added = bool(data.get("added")) + # the key for kid was changed/updated (optional) + updated = bool(data.get("updated")) + + return added or updated + + def add_keys(self, service: str, kid_keys: dict[Union[UUID, str], str]) -> int: + data = self.session.post( + url=f"{self.uri}/{service.lower()}", + json={"content_keys": {str(kid).replace("-", ""): key for kid, key in kid_keys.items()}}, + headers={"Accept": "application/json"}, + ).json() + + code = int(data.get("code", 0)) + message = data.get("message") + error = { + 0: None, + 1: Exceptions.AuthRejected, + 2: Exceptions.TooManyRequests, + 3: Exceptions.ServiceTagInvalid, + 4: Exceptions.KeyIdInvalid, + 5: Exceptions.ContentKeyInvalid, + }.get(code, ValueError) + + if error: + raise error(f"{message} ({code})") + + # each kid:key that was new to the vault (optional) + added = int(data.get("added")) + # each key for a kid that was changed/updated (optional) + updated = int(data.get("updated")) + + return added + updated + + def get_services(self) -> Iterator[str]: + data = self.session.post(url=self.uri, headers={"Accept": "application/json"}).json() + + code = int(data.get("code", 0)) + message = data.get("message") + error = { + 0: None, + 1: Exceptions.AuthRejected, + 2: Exceptions.TooManyRequests, + }.get(code, ValueError) + + if error: + raise error(f"{message} ({code})") + + service_list = data.get("service_list", []) + + if not isinstance(service_list, list): + raise ValueError(f"Expected {service_list} to be {list}, was {type(service_list)}") + + for service in service_list: + yield service + + +class Exceptions: + class AuthRejected(Exception): + """Authentication Error Occurred, is your token valid? Do you have permission to make this call?""" + + class TooManyRequests(Exception): + """Rate Limited; Sent too many requests in a given amount of time.""" + + class PageInvalid(Exception): + """Requested page does not exist.""" + + class ServiceTagInvalid(Exception): + """The Service Tag is invalid.""" + + class KeyIdInvalid(Exception): + """The Key ID is invalid.""" + + class ContentKeyInvalid(Exception): + """The Content Key is invalid.""" diff --git a/packages/envied/src/envied/vaults/HTTP.py b/packages/envied/src/envied/vaults/HTTP.py new file mode 100644 index 0000000..2fd171b --- /dev/null +++ b/packages/envied/src/envied/vaults/HTTP.py @@ -0,0 +1,337 @@ +import json +from enum import Enum +from typing import Iterator, Optional, Union +from uuid import UUID + +from requests import Session + +from envied.core import __version__ +from envied.core.vault import Vault + + +class InsertResult(Enum): + FAILURE = 0 + SUCCESS = 1 + ALREADY_EXISTS = 2 + + +class HTTP(Vault): + """Key Vault using HTTP API with support for both query parameters and JSON payloads.""" + + def __init__( + self, + name: str, + host: str, + password: str, + username: Optional[str] = None, + api_mode: str = "query", + no_push: bool = False, + ): + """ + Initialize HTTP Vault. + + Args: + name: Vault name + host: Host URL + password: Password for query mode or API token for json mode + username: Username (required for query mode, ignored for json mode) + api_mode: "query" for query parameters or "json" for JSON API + no_push: If True, this vault will not receive pushed keys + """ + super().__init__(name, no_push) + self.url = host + self.password = password + self.username = username + self.api_mode = api_mode.lower() + self.current_title = None + self.session = Session() + self.session.headers.update({"User-Agent": f"envied v{__version__}"}) + self.api_session_id = None + + # Validate configuration based on mode + if self.api_mode == "query" and not self.username: + raise ValueError("Username is required for query mode") + elif self.api_mode not in ["query", "json"]: + raise ValueError("api_mode must be either 'query' or 'json'") + + def request(self, method: str, params: dict = None) -> dict: + """Make a request to the JSON API vault.""" + if self.api_mode != "json": + raise ValueError("request method is only available in json mode") + + request_payload = { + "method": method, + "params": { + **(params or {}), + "session_id": self.api_session_id, + }, + "token": self.password, + } + + r = self.session.post(self.url, json=request_payload) + + if r.status_code == 404: + return {"status": "not_found"} + + if not r.ok: + raise ValueError(f"API returned HTTP Error {r.status_code}: {r.reason.title()}") + + try: + res = r.json() + except json.JSONDecodeError: + if r.status_code == 404: + return {"status": "not_found"} + raise ValueError(f"API returned an invalid response: {r.text}") + + if res.get("status_code") != 200: + raise ValueError(f"API returned an error: {res['status_code']} - {res['message']}") + + if session_id := res.get("message", {}).get("session_id"): + self.api_session_id = session_id + + return res.get("message", res) + + def get_key(self, kid: Union[UUID, str], service: str) -> Optional[str]: + if isinstance(kid, UUID): + kid = kid.hex + + if self.api_mode == "json": + try: + params = { + "kid": kid, + "service": service.lower(), + } + + response = self.request("GetKey", params) + if response.get("status") == "not_found": + return None + keys = response.get("keys", []) + for key_entry in keys: + if isinstance(key_entry, str) and ":" in key_entry: + entry_kid, entry_key = key_entry.split(":", 1) + if entry_kid == kid: + return entry_key + elif isinstance(key_entry, dict): + if key_entry.get("kid") == kid: + return key_entry.get("key") + except Exception as e: + print(f"Failed to get key ({e.__class__.__name__}: {e})") + return None + return None + else: # query mode + response = self.session.get( + self.url, + params={"service": service.lower(), "username": self.username, "password": self.password, "kid": kid}, + ) + + data = response.json() + + if data.get("status_code") != 200 or not data.get("keys"): + return None + + return data["keys"][0]["key"] + + def get_keys(self, service: str) -> Iterator[tuple[str, str]]: + if self.api_mode == "json": + # JSON API doesn't support getting all keys, so return empty iterator + # This will cause the copy command to rely on the API's internal duplicate handling + return iter([]) + else: # query mode + response = self.session.get( + self.url, params={"service": service.lower(), "username": self.username, "password": self.password} + ) + + data = response.json() + + if data.get("status_code") != 200 or not data.get("keys"): + return + + for key_entry in data["keys"]: + yield key_entry["kid"], key_entry["key"] + + def add_key(self, service: str, kid: Union[UUID, str], key: str) -> bool: + if not key or key.count("0") == len(key): + raise ValueError("You cannot add a NULL Content Key to a Vault.") + + if isinstance(kid, UUID): + kid = kid.hex + + title = getattr(self, "current_title", None) + + if self.api_mode == "json": + try: + response = self.request( + "InsertKey", + { + "kid": kid, + "key": key, + "service": service.lower(), + "title": title, + }, + ) + if response.get("status") == "not_found": + return False + return response.get("inserted", False) + except Exception: + return False + else: # query mode + response = self.session.get( + self.url, + params={ + "service": service.lower(), + "username": self.username, + "password": self.password, + "kid": kid, + "key": key, + "title": title, + }, + ) + + data = response.json() + + return data.get("status_code") == 200 + + def add_keys(self, service: str, kid_keys: dict[Union[UUID, str], str]) -> int: + for kid, key in kid_keys.items(): + if not key or key.count("0") == len(key): + raise ValueError("You cannot add a NULL Content Key to a Vault.") + + processed_kid_keys = { + str(kid).replace("-", "") if isinstance(kid, UUID) else kid: key for kid, key in kid_keys.items() + } + + inserted_count = 0 + title = getattr(self, "current_title", None) + + if self.api_mode == "json": + for kid, key in processed_kid_keys.items(): + try: + response = self.request( + "InsertKey", + { + "kid": kid, + "key": key, + "service": service.lower(), + "title": title, + }, + ) + if response.get("status") == "not_found": + continue + if response.get("inserted", False): + inserted_count += 1 + except Exception: + continue + else: # query mode + for kid, key in processed_kid_keys.items(): + response = self.session.get( + self.url, + params={ + "service": service.lower(), + "username": self.username, + "password": self.password, + "kid": kid, + "key": key, + "title": title, + }, + ) + + data = response.json() + + if data.get("status_code") == 200 and data.get("inserted", True): + inserted_count += 1 + + return inserted_count + + def get_services(self) -> Iterator[str]: + if self.api_mode == "json": + try: + response = self.request("GetServices") + services = response.get("services", []) + for service in services: + yield service + except Exception: + return iter([]) + else: # query mode + response = self.session.get( + self.url, params={"username": self.username, "password": self.password, "list_services": True} + ) + + data = response.json() + + if data.get("status_code") != 200: + return + + services = data.get("services", []) + for service in services: + yield service + + def set_title(self, title: str): + """ + Set a title to be used for the next key insertions. + This is optional and will be sent with add_key requests if available. + """ + self.current_title = title + + def insert_key_with_result( + self, service: str, kid: Union[UUID, str], key: str, title: Optional[str] = None + ) -> InsertResult: + """ + Insert a key and return detailed result information. + This method provides more granular feedback than the standard add_key method. + Available in both API modes. + """ + if not key or key.count("0") == len(key): + raise ValueError("You cannot add a NULL Content Key to a Vault.") + + if isinstance(kid, UUID): + kid = kid.hex + + if title is None: + title = getattr(self, "current_title", None) + + if self.api_mode == "json": + try: + response = self.request( + "InsertKey", + { + "kid": kid, + "key": key, + "service": service.lower(), + "title": title, + }, + ) + + if response.get("status") == "not_found": + return InsertResult.FAILURE + + if response.get("inserted", False): + return InsertResult.SUCCESS + else: + return InsertResult.ALREADY_EXISTS + + except Exception: + return InsertResult.FAILURE + else: # query mode + response = self.session.get( + self.url, + params={ + "service": service.lower(), + "username": self.username, + "password": self.password, + "kid": kid, + "key": key, + "title": title, + }, + ) + + try: + data = response.json() + if data.get("status_code") == 200: + if data.get("inserted", True): + return InsertResult.SUCCESS + else: + return InsertResult.ALREADY_EXISTS + else: + return InsertResult.FAILURE + except Exception: + return InsertResult.FAILURE diff --git a/packages/envied/src/envied/vaults/HTTPAPI.py b/packages/envied/src/envied/vaults/HTTPAPI.py new file mode 100644 index 0000000..5836954 --- /dev/null +++ b/packages/envied/src/envied/vaults/HTTPAPI.py @@ -0,0 +1,215 @@ +import json +from typing import Iterator, Optional, Union +from uuid import UUID +from enum import Enum + +from requests import Session + +from envied.core import __version__ +from envied.core.vault import Vault + + +class InsertResult(Enum): + FAILURE = 0 + SUCCESS = 1 + ALREADY_EXISTS = 2 + + +class HTTPAPI(Vault): + """Key Vault using a structured HTTP API with JSON payloads and token authentication.""" + + def __init__(self, name: str, host: str, password: str): + super().__init__(name) + self.url = host + self.password = password # This is the API token + self.current_title = None + self.session = Session() + self.session.headers.update({"User-Agent": f"Devine v{__version__}"}) + self.api_session_id = None + + def request(self, method: str, params: dict = None) -> dict: + """Make a request to the HTTPAPI vault.""" + request_payload = { + "method": method, + "params": { + **(params or {}), + "session_id": self.api_session_id, + }, + "token": self.password, + } + + r = self.session.post(self.url, json=request_payload) + + if r.status_code == 404: + return {"status": "not_found"} + + if not r.ok: + raise ValueError(f"API returned HTTP Error {r.status_code}: {r.reason.title()}") + + try: + res = r.json() + except json.JSONDecodeError: + if r.status_code == 404: + return {"status": "not_found"} + raise ValueError(f"API returned an invalid response: {r.text}") + + if res.get("status_code") != 200: + raise ValueError(f"API returned an error: {res['status_code']} - {res['message']}") + + if session_id := res.get("message", {}).get("session_id"): + self.api_session_id = session_id + + return res.get("message", res) + + def get_key(self, kid: Union[UUID, str], service: str) -> Optional[str]: + if isinstance(kid, UUID): + kid = kid.hex + + try: + title = getattr(self, "current_title", None) + response = self.request( + "GetKey", + { + "kid": kid, + "service": service.lower(), + "title": title, + }, + ) + if response.get("status") == "not_found": + return None + keys = response.get("keys", []) + for key_entry in keys: + if key_entry["kid"] == kid: + return key_entry["key"] + except Exception as e: + print(f"Failed to get key ({e.__class__.__name__}: {e})") + return None + + return None + + def get_keys(self, service: str) -> Iterator[tuple[str, str]]: + """Get all keys for a service - this may need to be implemented based on your API.""" + try: + response = self.request( + "GetKeys", + { + "service": service.lower(), + }, + ) + keys = response.get("keys", []) + for key_entry in keys: + yield key_entry["kid"], key_entry["key"] + except Exception: + return iter([]) + + def add_key(self, service: str, kid: Union[UUID, str], key: str) -> bool: + if not key or key.count("0") == len(key): + raise ValueError("You cannot add a NULL Content Key to a Vault.") + + if isinstance(kid, UUID): + kid = kid.hex + + title = getattr(self, "current_title", None) + + try: + response = self.request( + "InsertKey", + { + "kid": kid, + "key": key, + "service": service.lower(), + "title": title, + }, + ) + if response.get("status") == "not_found": + return False + return response.get("inserted", False) + except Exception: + return False + + def add_keys(self, service: str, kid_keys: dict[Union[UUID, str], str]) -> int: + for kid, key in kid_keys.items(): + if not key or key.count("0") == len(key): + raise ValueError("You cannot add a NULL Content Key to a Vault.") + + processed_kid_keys = { + str(kid).replace("-", "") if isinstance(kid, UUID) else kid: key for kid, key in kid_keys.items() + } + + inserted_count = 0 + title = getattr(self, "current_title", None) + + for kid, key in processed_kid_keys.items(): + try: + response = self.request( + "InsertKey", + { + "kid": kid, + "key": key, + "service": service.lower(), + "title": title, + }, + ) + if response.get("status") == "not_found": + continue + if response.get("inserted", False): + inserted_count += 1 + except Exception: + continue + + return inserted_count + + def get_services(self) -> Iterator[str]: + """Get available services - this may need to be implemented based on your API.""" + try: + response = self.request("GetServices") + services = response.get("services", []) + for service in services: + yield service + except Exception: + return iter([]) + + def set_title(self, title: str): + """ + Set a title to be used for the next key insertions. + This is optional and will be sent with add_key requests if available. + """ + self.current_title = title + + def insert_key_with_result( + self, service: str, kid: Union[UUID, str], key: str, title: Optional[str] = None + ) -> InsertResult: + """ + Insert a key and return detailed result information. + This method provides more granular feedback than the standard add_key method. + """ + if not key or key.count("0") == len(key): + raise ValueError("You cannot add a NULL Content Key to a Vault.") + + if isinstance(kid, UUID): + kid = kid.hex + + if title is None: + title = getattr(self, "current_title", None) + + try: + response = self.request( + "InsertKey", + { + "kid": kid, + "key": key, + "service": service.lower(), + "title": title, + }, + ) + + if response.get("status") == "not_found": + return InsertResult.FAILURE + + if response.get("inserted", False): + return InsertResult.SUCCESS + else: + return InsertResult.ALREADY_EXISTS + + except Exception: + return InsertResult.FAILURE diff --git a/packages/envied/src/envied/vaults/MySQL.py b/packages/envied/src/envied/vaults/MySQL.py new file mode 100644 index 0000000..e9c33c8 --- /dev/null +++ b/packages/envied/src/envied/vaults/MySQL.py @@ -0,0 +1,244 @@ +import threading +from typing import Iterator, Optional, Union +from uuid import UUID + +import pymysql +from pymysql.cursors import DictCursor + +from envied.core.services import Services +from envied.core.vault import Vault + + +class MySQL(Vault): + """Key Vault using a remotely-accessed mysql database connection.""" + + def __init__(self, name: str, host: str, database: str, username: str, no_push: bool = False, **kwargs): + """ + All extra arguments provided via **kwargs will be sent to pymysql.connect. + This can be used to provide more specific connection information. + """ + super().__init__(name, no_push) + self.slug = f"{host}:{database}:{username}" + self.conn_factory = ConnectionFactory( + dict(host=host, db=database, user=username, cursorclass=DictCursor, **kwargs) + ) + + self.permissions = self.get_permissions() + if not self.has_permission("SELECT"): + raise PermissionError(f"MySQL vault {self.slug} has no SELECT permission.") + + def get_key(self, kid: Union[UUID, str], service: str) -> Optional[str]: + if not self.has_table(service): + # no table, no key, simple + return None + + if isinstance(kid, UUID): + kid = kid.hex + + conn = self.conn_factory.get() + cursor = conn.cursor() + + try: + cursor.execute( + # TODO: SQL injection risk + f"SELECT `id`, `key_` FROM `{service}` WHERE `kid`=%s AND `key_`!=%s", + (kid, "0" * 32), + ) + cek = cursor.fetchone() + if not cek: + return None + return cek["key_"] + finally: + cursor.close() + + def get_keys(self, service: str) -> Iterator[tuple[str, str]]: + if not self.has_table(service): + # no table, no keys, simple + return None + + conn = self.conn_factory.get() + cursor = conn.cursor() + + try: + cursor.execute( + # TODO: SQL injection risk + f"SELECT `kid`, `key_` FROM `{service}` WHERE `key_`!=%s", + ("0" * 32,), + ) + for row in cursor.fetchall(): + yield row["kid"], row["key_"] + finally: + cursor.close() + + def add_key(self, service: str, kid: Union[UUID, str], key: str) -> bool: + if not key or key.count("0") == len(key): + raise ValueError("You cannot add a NULL Content Key to a Vault.") + + if not self.has_permission("INSERT", table=service): + raise PermissionError(f"MySQL vault {self.slug} has no INSERT permission.") + + if not self.has_table(service): + try: + self.create_table(service) + except PermissionError: + return False + + if isinstance(kid, UUID): + kid = kid.hex + + conn = self.conn_factory.get() + cursor = conn.cursor() + + try: + cursor.execute( + # TODO: SQL injection risk + f"SELECT `id` FROM `{service}` WHERE `kid`=%s AND `key_`=%s", + (kid, key), + ) + if cursor.fetchone(): + # table already has this exact KID:KEY stored + return True + cursor.execute( + # TODO: SQL injection risk + f"INSERT INTO `{service}` (kid, key_) VALUES (%s, %s)", + (kid, key), + ) + finally: + conn.commit() + cursor.close() + + return True + + def add_keys(self, service: str, kid_keys: dict[Union[UUID, str], str]) -> int: + for kid, key in kid_keys.items(): + if not key or key.count("0") == len(key): + raise ValueError("You cannot add a NULL Content Key to a Vault.") + + if not self.has_permission("INSERT", table=service): + raise PermissionError(f"MySQL vault {self.slug} has no INSERT permission.") + + if not self.has_table(service): + try: + self.create_table(service) + except PermissionError: + return 0 + + if not isinstance(kid_keys, dict): + raise ValueError(f"The kid_keys provided is not a dictionary, {kid_keys!r}") + if not all(isinstance(kid, (str, UUID)) and isinstance(key_, str) for kid, key_ in kid_keys.items()): + raise ValueError("Expecting dict with Key of str/UUID and value of str.") + + if any(isinstance(kid, UUID) for kid, key_ in kid_keys.items()): + kid_keys = {kid.hex if isinstance(kid, UUID) else kid: key_ for kid, key_ in kid_keys.items()} + + conn = self.conn_factory.get() + cursor = conn.cursor() + + try: + cursor.executemany( + # TODO: SQL injection risk + f"INSERT IGNORE INTO `{service}` (kid, key_) VALUES (%s, %s)", + kid_keys.items(), + ) + return cursor.rowcount + finally: + conn.commit() + cursor.close() + + def get_services(self) -> Iterator[str]: + conn = self.conn_factory.get() + cursor = conn.cursor() + + try: + cursor.execute("SHOW TABLES") + for table in cursor.fetchall(): + # each entry has a key named `Tables_in_` + yield Services.get_tag(list(table.values())[0]) + finally: + cursor.close() + + def has_table(self, name: str) -> bool: + """Check if the Vault has a Table with the specified name.""" + conn = self.conn_factory.get() + cursor = conn.cursor() + + try: + cursor.execute( + "SELECT count(TABLE_NAME) FROM information_schema.TABLES WHERE TABLE_SCHEMA=%s AND TABLE_NAME=%s", + (conn.db, name), + ) + return list(cursor.fetchone().values())[0] == 1 + finally: + cursor.close() + + def create_table(self, name: str): + """Create a Table with the specified name if not yet created.""" + if self.has_table(name): + return + + if not self.has_permission("CREATE"): + raise PermissionError(f"MySQL vault {self.slug} has no CREATE permission.") + + conn = self.conn_factory.get() + cursor = conn.cursor() + + try: + cursor.execute( + # TODO: SQL injection risk + f""" + CREATE TABLE IF NOT EXISTS {name} ( + id int AUTO_INCREMENT PRIMARY KEY, + kid VARCHAR(64) NOT NULL, + key_ VARCHAR(64) NOT NULL, + UNIQUE(kid, key_) + ); + """ + ) + finally: + conn.commit() + cursor.close() + + def get_permissions(self) -> list: + """Get and parse Grants to a more easily usable list tuple array.""" + conn = self.conn_factory.get() + cursor = conn.cursor() + + try: + cursor.execute("SHOW GRANTS") + grants = cursor.fetchall() + grants = [next(iter(x.values())) for x in grants] + grants = [tuple(x[6:].split(" TO ")[0].split(" ON ")) for x in list(grants)] + grants = [ + ( + list(map(str.strip, perms.replace("ALL PRIVILEGES", "*").split(","))), + location.replace("`", "").split("."), + ) + for perms, location in grants + ] + return grants + finally: + conn.commit() + cursor.close() + + def has_permission(self, operation: str, database: Optional[str] = None, table: Optional[str] = None) -> bool: + """Check if the current connection has a specific permission.""" + grants = [x for x in self.permissions if x[0] == ["*"] or operation.upper() in x[0]] + if grants and database: + grants = [x for x in grants if x[1][0] in (database, "*")] + if grants and table: + grants = [x for x in grants if x[1][1] in (table, "*")] + return bool(grants) + + +class ConnectionFactory: + def __init__(self, con: dict): + self._con = con + self._store = threading.local() + + def _create_connection(self) -> pymysql.Connection: + return pymysql.connect(**self._con) + + def get(self) -> pymysql.Connection: + if not hasattr(self._store, "conn"): + self._store.conn = self._create_connection() + return self._store.conn diff --git a/packages/envied/src/envied/vaults/SQLite.py b/packages/envied/src/envied/vaults/SQLite.py new file mode 100644 index 0000000..2ab50f7 --- /dev/null +++ b/packages/envied/src/envied/vaults/SQLite.py @@ -0,0 +1,179 @@ +import sqlite3 +import threading +from pathlib import Path +from sqlite3 import Connection +from typing import Iterator, Optional, Union +from uuid import UUID + +from envied.core.services import Services +from envied.core.vault import Vault + + +class SQLite(Vault): + """Key Vault using a locally-accessed sqlite DB file.""" + + def __init__(self, name: str, path: Union[str, Path], no_push: bool = False): + super().__init__(name, no_push) + self.path = Path(path).expanduser() + # TODO: Use a DictCursor or such to get fetches as dict? + self.conn_factory = ConnectionFactory(self.path) + + def get_key(self, kid: Union[UUID, str], service: str) -> Optional[str]: + if not self.has_table(service): + # no table, no key, simple + return None + + if isinstance(kid, UUID): + kid = kid.hex + + conn = self.conn_factory.get() + cursor = conn.cursor() + + try: + cursor.execute(f"SELECT `id`, `key_` FROM `{service}` WHERE `kid`=? AND `key_`!=?", (kid, "0" * 32)) + cek = cursor.fetchone() + if not cek: + return None + return cek[1] + finally: + cursor.close() + + def get_keys(self, service: str) -> Iterator[tuple[str, str]]: + if not self.has_table(service): + # no table, no keys, simple + return None + + conn = self.conn_factory.get() + cursor = conn.cursor() + + try: + cursor.execute(f"SELECT `kid`, `key_` FROM `{service}` WHERE `key_`!=?", ("0" * 32,)) + for kid, key_ in cursor.fetchall(): + yield kid, key_ + finally: + cursor.close() + + def add_key(self, service: str, kid: Union[UUID, str], key: str) -> bool: + if not key or key.count("0") == len(key): + raise ValueError("You cannot add a NULL Content Key to a Vault.") + + if not self.has_table(service): + self.create_table(service) + + if isinstance(kid, UUID): + kid = kid.hex + + conn = self.conn_factory.get() + cursor = conn.cursor() + + try: + cursor.execute( + # TODO: SQL injection risk + f"SELECT `id` FROM `{service}` WHERE `kid`=? AND `key_`=?", + (kid, key), + ) + if cursor.fetchone(): + # table already has this exact KID:KEY stored + return True + cursor.execute( + # TODO: SQL injection risk + f"INSERT INTO `{service}` (kid, key_) VALUES (?, ?)", + (kid, key), + ) + finally: + conn.commit() + cursor.close() + + return True + + def add_keys(self, service: str, kid_keys: dict[Union[UUID, str], str]) -> int: + for kid, key in kid_keys.items(): + if not key or key.count("0") == len(key): + raise ValueError("You cannot add a NULL Content Key to a Vault.") + + if not self.has_table(service): + self.create_table(service) + + if not isinstance(kid_keys, dict): + raise ValueError(f"The kid_keys provided is not a dictionary, {kid_keys!r}") + if not all(isinstance(kid, (str, UUID)) and isinstance(key_, str) for kid, key_ in kid_keys.items()): + raise ValueError("Expecting dict with Key of str/UUID and value of str.") + + if any(isinstance(kid, UUID) for kid, key_ in kid_keys.items()): + kid_keys = {kid.hex if isinstance(kid, UUID) else kid: key_ for kid, key_ in kid_keys.items()} + + conn = self.conn_factory.get() + cursor = conn.cursor() + + try: + cursor.executemany( + # TODO: SQL injection risk + f"INSERT OR IGNORE INTO `{service}` (kid, key_) VALUES (?, ?)", + kid_keys.items(), + ) + return cursor.rowcount + finally: + conn.commit() + cursor.close() + + def get_services(self) -> Iterator[str]: + conn = self.conn_factory.get() + cursor = conn.cursor() + + try: + cursor.execute("SELECT name FROM sqlite_master WHERE type='table';") + for (name,) in cursor.fetchall(): + if name != "sqlite_sequence": + yield Services.get_tag(name) + finally: + cursor.close() + + def has_table(self, name: str) -> bool: + """Check if the Vault has a Table with the specified name.""" + conn = self.conn_factory.get() + cursor = conn.cursor() + + try: + cursor.execute("SELECT count(name) FROM sqlite_master WHERE type='table' AND name=?", (name,)) + return cursor.fetchone()[0] == 1 + finally: + cursor.close() + + def create_table(self, name: str): + """Create a Table with the specified name if not yet created.""" + if self.has_table(name): + return + + conn = self.conn_factory.get() + cursor = conn.cursor() + + try: + cursor.execute( + # TODO: SQL injection risk + f""" + CREATE TABLE IF NOT EXISTS {name} ( + "id" INTEGER NOT NULL UNIQUE, + "kid" TEXT NOT NULL COLLATE NOCASE, + "key_" TEXT NOT NULL COLLATE NOCASE, + PRIMARY KEY("id" AUTOINCREMENT), + UNIQUE("kid", "key_") + ); + """ + ) + finally: + conn.commit() + cursor.close() + + +class ConnectionFactory: + def __init__(self, path: Union[str, Path]): + self._path = path + self._store = threading.local() + + def _create_connection(self) -> Connection: + return sqlite3.connect(self._path) + + def get(self) -> Connection: + if not hasattr(self._store, "conn"): + self._store.conn = self._create_connection() + return self._store.conn diff --git a/packages/envied/src/envied/vaults/__init__.py b/packages/envied/src/envied/vaults/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/packages/vinefeeder/src/vinefeeder/__init__.py b/packages/vinefeeder/src/vinefeeder/__init__.py index 0058b93..bb35ee1 100644 --- a/packages/vinefeeder/src/vinefeeder/__init__.py +++ b/packages/vinefeeder/src/vinefeeder/__init__.py @@ -1 +1 @@ -__version__ = "1.0.1" \ No newline at end of file +__version__ = "1.0.2" \ No newline at end of file diff --git a/packages/vinefeeder/src/vinefeeder/base_loader.py b/packages/vinefeeder/src/vinefeeder/base_loader.py index a6e1a27..a283fcd 100644 --- a/packages/vinefeeder/src/vinefeeder/base_loader.py +++ b/packages/vinefeeder/src/vinefeeder/base_loader.py @@ -406,7 +406,7 @@ class BaseLoader: if self.TERMINAL_RESET and not self.BATCH_DOWNLOAD: console.print(f"[{catppuccin_mocha['text2']}]Preparing to reset Terminal[/]") time.sleep(5) - self.reset_terminal() + #self.reset_terminal() console.print(f"[{catppuccin_mocha['pink']}]Ready![/]")