diff --git a/unshackle/commands/dl.py b/unshackle/commands/dl.py index f3bfef2e..ec218667 100644 --- a/unshackle/commands/dl.py +++ b/unshackle/commands/dl.py @@ -4620,7 +4620,7 @@ class dl: cek_tree = pre_existing_tree for kid_, key in drm.content_keys.items(): - label = f"[text2]{kid_.hex}:{key}" + label = f"[text2]{kid_.hex}:{key}*" if not any(f"{kid_.hex}:{key}" in x.label for x in cek_tree.children): cek_tree.add(label) diff --git a/unshackle/commands/env.py b/unshackle/commands/env.py index 7813ad32..152272c3 100644 --- a/unshackle/commands/env.py +++ b/unshackle/commands/env.py @@ -19,10 +19,28 @@ from unshackle.core.temp import TASK_PREFIX, is_stale def get_dependencies() -> list[dict]: """Binary dependency inventory shared by `env check` and the API env check.""" - return [ - {"name": "FFmpeg", "binary": binaries.FFMPEG, "required": True, "desc": "Media processing", "cat": "Core"}, - {"name": "FFprobe", "binary": binaries.FFProbe, "required": True, "desc": "Media analysis", "cat": "Core"}, - {"name": "MKVToolNix", "binary": binaries.MKVToolNix, "required": True, "desc": "MKV muxing", "cat": "Core"}, + deps = [ + { + "name": "FFmpeg", + "binary": binaries.FFMPEG, + "required": True, + "desc": "Media processing", + "cat": "Core", + }, + { + "name": "FFprobe", + "binary": binaries.FFProbe, + "required": True, + "desc": "Media analysis", + "cat": "Core", + }, + { + "name": "MKVToolNix", + "binary": binaries.MKVToolNix, + "required": True, + "desc": "MKV muxing", + "cat": "Core", + }, { "name": "mkvpropedit", "binary": binaries.Mkvpropedit, @@ -44,7 +62,13 @@ def get_dependencies() -> list[dict]: "desc": "DRM decryption", "cat": "DRM", }, - {"name": "dovi_tool", "binary": binaries.DoviTool, "required": False, "desc": "Dolby Vision", "cat": "HDR"}, + { + "name": "dovi_tool", + "binary": binaries.DoviTool, + "required": False, + "desc": "Dolby Vision", + "cat": "HDR", + }, { "name": "HDR10Plus_tool", "binary": binaries.HDR10PlusTool, @@ -66,8 +90,20 @@ def get_dependencies() -> list[dict]: "desc": "CC extraction", "cat": "Subtitle", }, - {"name": "FFplay", "binary": binaries.FFPlay, "required": False, "desc": "Simple player", "cat": "Player"}, - {"name": "MPV", "binary": binaries.MPV, "required": False, "desc": "Advanced player", "cat": "Player"}, + { + "name": "FFplay", + "binary": binaries.FFPlay, + "required": False, + "desc": "Simple player", + "cat": "Player", + }, + { + "name": "MPV", + "binary": binaries.MPV, + "required": False, + "desc": "Advanced player", + "cat": "Player", + }, { "name": "HolaProxy", "binary": binaries.HolaProxy, @@ -75,11 +111,42 @@ def get_dependencies() -> list[dict]: "desc": "Proxy service", "cat": "Network", }, - {"name": "Caddy", "binary": binaries.Caddy, "required": False, "desc": "Web server", "cat": "Network"}, - {"name": "Docker", "binary": binaries.Docker, "required": False, "desc": "Gluetun VPN", "cat": "Network"}, - {"name": "git", "binary": binaries.Git, "required": False, "desc": "Service repos", "cat": "Network"}, + { + "name": "Caddy", + "binary": binaries.Caddy, + "required": False, + "desc": "Web server", + "cat": "Network", + }, + { + "name": "Docker", + "binary": binaries.Docker, + "required": False, + "desc": "Gluetun VPN", + "cat": "Network", + }, + { + "name": "git", + "binary": binaries.Git, + "required": False, + "desc": "Service repos", + "cat": "Network", + }, ] + for reg in binaries.get_registered_dependencies(): + deps.append( + { + "name": reg["name"], + "binary": getattr(binaries, reg["attr"], None), + "required": False, + "desc": reg["desc"], + "cat": "Service", + } + ) + + return deps + def clear_directory(path: Path) -> tuple[int, int]: """Delete a directory's contents, returning (files_removed, freed_bytes). Recreates the dir. diff --git a/unshackle/core/binaries.py b/unshackle/core/binaries.py index daa77506..30f41d28 100644 --- a/unshackle/core/binaries.py +++ b/unshackle/core/binaries.py @@ -6,18 +6,31 @@ from typing import Optional __shaka_platform = {"win32": "win", "darwin": "osx"}.get(sys.platform, sys.platform) -def find(*names: str) -> Optional[Path]: +def find(*names: str, search_dirs: Optional[list[Path]] = None) -> Optional[Path]: """Find the path of the first found binary name.""" current_dir = Path(__file__).resolve().parent.parent local_binaries_dir = current_dir / "binaries" + services_dir = current_dir / "services" + + dirs_to_check: list[Path] = [local_binaries_dir] + if services_dir.exists(): + for s_dir in services_dir.iterdir(): + s_bin = s_dir / "binaries" + if s_bin.is_dir(): + dirs_to_check.append(s_bin) + + if search_dirs: + dirs_to_check.extend(search_dirs) ext = ".exe" if sys.platform == "win32" else "" for name in names: - if local_binaries_dir.exists(): - candidate_paths = [local_binaries_dir / f"{name}{ext}", local_binaries_dir / name / f"{name}{ext}"] + for base_dir in dirs_to_check: + if not base_dir.exists(): + continue + candidate_paths = [base_dir / f"{name}{ext}", base_dir / name / f"{name}{ext}"] - for subdir in local_binaries_dir.iterdir(): + for subdir in base_dir.iterdir(): if subdir.is_dir(): candidate_paths.append(subdir / f"{name}{ext}") @@ -59,15 +72,53 @@ __binaries = { "HDR10PlusTool": ("hdr10plus_tool", "HDR10Plus_tool"), "Mp4decrypt": ("mp4decrypt",), "Docker": ("docker",), - "ML_Worker": ("ML-Worker",), "Git": ("git",), } +_registered_binaries: dict[str, dict] = {} + + +def register( + attr_name: str, + *candidates: str, + desc: str = "Custom tool", +) -> None: + """Register a custom binary dynamically.""" + file_candidates = candidates if candidates else (attr_name.lower(),) + __binaries[attr_name] = file_candidates + _registered_binaries[attr_name] = { + "name": attr_name, + "attr": attr_name, + "desc": desc, + } + + +def get_registered_dependencies() -> list[dict]: + """Return all registered custom binary dependencies for env check.""" + return list(_registered_binaries.values()) + + +_services_loaded = False + def __getattr__(name: str) -> Optional[Path]: + global _services_loaded candidates = __binaries.get(name) + if candidates is None and not _services_loaded: + _services_loaded = True + try: + candidates = __binaries.get(name) + except Exception: + pass + if candidates is None: + # Fallback to direct name and lowercase candidate + resolved = find(name, name.lower()) + if resolved is not None: + globals()[name] = resolved + return resolved raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + # Cache the result (including None for an absent binary) so later access is a plain attr hit. resolved = find(*candidates) globals()[name] = resolved @@ -90,7 +141,8 @@ __all__ = ( "HDR10PlusTool", "Mp4decrypt", "Docker", - "ML_Worker", "Git", "find", + "register", + "get_registered_dependencies", ) diff --git a/unshackle/core/cdm/monalisa/monalisa_cdm.py b/unshackle/core/cdm/monalisa/monalisa_cdm.py index 548e6e51..e1561759 100644 --- a/unshackle/core/cdm/monalisa/monalisa_cdm.py +++ b/unshackle/core/cdm/monalisa/monalisa_cdm.py @@ -14,12 +14,10 @@ import re import sys import uuid from pathlib import Path -from typing import Dict, Optional, Union +from typing import Dict, Union import wasmtime -from unshackle.core import binaries - logger = logging.getLogger(__name__) @@ -93,13 +91,6 @@ class MonaLisaCDM: self.exports = {} self.ctx = None - @staticmethod - def get_worker_path() -> Optional[Path]: - """Get ML-Worker binary path from the unshackle binaries system.""" - if binaries.ML_Worker: - return Path(binaries.ML_Worker) - return None - def open(self) -> int: """ Open a CDM session. diff --git a/unshackle/core/drm/monalisa.py b/unshackle/core/drm/monalisa.py index afcd4a58..22ed4463 100644 --- a/unshackle/core/drm/monalisa.py +++ b/unshackle/core/drm/monalisa.py @@ -1,22 +1,20 @@ """ MonaLisa DRM System. -A WASM-based DRM system that uses local content key extraction and two-stage -segment decryption (ML-Worker binary + AES-ECB). +A WASM-based DRM system that uses local content key extraction from ticket data. """ from __future__ import annotations +import base64 import logging -import os import subprocess import sys from pathlib import Path from typing import Any, Optional, Union from uuid import UUID -from Cryptodome.Cipher import AES -from Cryptodome.Util.Padding import unpad +from unshackle.core.cdm.monalisa import MonaLisaCDM log = logging.getLogger(__name__) @@ -28,10 +26,6 @@ class MonaLisa: Unlike Widevine/PlayReady, MonaLisa does not use a challenge/response flow with a license server. Instead, the service API gives the PSSH value (ticket) directly, and a WASM module extracts the content keys locally. - - Decryption has two stages: - 1. ML-Worker binary: Removes MonaLisa encryption layer (bbts -> ents) - 2. AES-ECB decryption: Final decryption with the service-supplied content key """ class Exceptions: @@ -41,17 +35,11 @@ class MonaLisa: class KeyExtractionFailed(Exception): """Raised when content key extraction from the ticket fails.""" - class WorkerNotFound(Exception): - """Raised when the ML-Worker binary is not found.""" - - class DecryptionFailed(Exception): - """Raised when segment decryption fails.""" - def __init__( self, ticket: Union[str, bytes], - aes_key: Union[str, bytes], device_path: Path, + tool_path: Optional[Path] = None, **kwargs: Any, ): """ @@ -59,8 +47,8 @@ class MonaLisa: Args: ticket: PSSH value from service API (base64 string or raw bytes). - aes_key: AES-ECB content key for second-stage decryption (hex string or bytes). device_path: Path to the CDM device file (.mld). + tool_path: Optional path to an external decryption binary. **kwargs: Additional metadata stored in self.data. Raises: @@ -71,13 +59,9 @@ class MonaLisa: raise MonaLisa.Exceptions.TicketNotFound("No PSSH/ticket data provided.") self._ticket = ticket - - if isinstance(aes_key, str): - self._aes_key = bytes.fromhex(aes_key) - else: - self._aes_key = aes_key - self._device_path = device_path + self._tool_path = tool_path + self._decrypted_segments: int = 0 self._kid: Optional[UUID] = None self._key: Optional[str] = None self.data: dict = kwargs or {} @@ -86,8 +70,6 @@ class MonaLisa: def extract_keys(self) -> None: """Extract keys from the ticket using the MonaLisa CDM.""" - # Import here to avoid circular import - from unshackle.core.cdm.monalisa import MonaLisaCDM try: cdm = MonaLisaCDM(device_path=self._device_path) @@ -108,21 +90,25 @@ class MonaLisa: def from_ticket( cls, ticket: Union[str, bytes], - aes_key: Union[str, bytes], device_path: Path, + tool_path: Optional[Path] = None, ) -> MonaLisa: """ Make a MonaLisa DRM instance from a PSSH/ticket. Args: ticket: PSSH value from service API. - aes_key: AES-ECB content key for second-stage decryption. device_path: Path to the CDM device file (.mld). + tool_path: Optional path to an external decryption binary. Returns: MonaLisa DRM instance with extracted keys. """ - return cls(ticket=ticket, aes_key=aes_key, device_path=device_path) + return cls( + ticket=ticket, + device_path=device_path, + tool_path=tool_path, + ) @property def kid(self) -> Optional[UUID]: @@ -166,8 +152,6 @@ class MonaLisa: Returns: The Content ID string if extractable, None otherwise. """ - import base64 - try: if isinstance(self._ticket, bytes): data = self._ticket @@ -196,47 +180,80 @@ class MonaLisa: return {self._kid: self._key} return {} - def decrypt_segment(self, segment_path: Path) -> None: - """ - Decrypt a single segment using two-stage decryption. + @property + def key_pair(self) -> str: + """Get formatted KID:KEY pair for CLI decryptors (32hex:32hex).""" + if self._kid and self._key: + return f"{self._kid.hex}:{self._key}" + return "" - Stage 1: ML-Worker binary (bbts -> ents) - Stage 2: AES-ECB decryption (ents -> ts) + def decrypt_segment(self, segment_path: Path, output_path: Optional[Path] = None) -> None: + """ + Decrypt an individual segment file using the configured CLI tool via subprocess. Args: - segment_path: Path to the encrypted segment file. - - Raises: - WorkerNotFound: If ML-Worker binary is not available. - DecryptionFailed: If decryption fails at any stage. + segment_path: Path to the segment file to decrypt. + output_path: Optional path for the decrypted output segment. """ - if not self._key: + self._execute_decrypt(segment_path, output_path) + self._decrypted_segments += 1 + + def decrypt(self, path: Path, output_path: Optional[Path] = None, *, is_segment: bool = False) -> None: + """ + Decrypt a media file (segment or track) using the configured CLI tool via subprocess. + + Args: + path: Path to the target file to decrypt. + output_path: Optional path for the decrypted output file. + is_segment: Set True if this is an individual segment rather than a full track. + """ + if is_segment: + return self.decrypt_segment(path, output_path) + + # If segments were already decrypted during chunk downloads, skip reprocessing on final track. + if self._decrypted_segments > 0: + log.debug( + "MonaLisa: Track already decrypted during segment downloads (%s segments)", self._decrypted_segments + ) return - # Import here to avoid circular import - from unshackle.core.cdm.monalisa import MonaLisaCDM + self._execute_decrypt(path, output_path) - worker_path = MonaLisaCDM.get_worker_path() - if not worker_path or not worker_path.exists(): - raise MonaLisa.Exceptions.WorkerNotFound("ML-Worker not found.") + def _execute_decrypt(self, path: Path, output_path: Optional[Path] = None) -> None: + """Execute external CLI tool via subprocess to decrypt the target file.""" + if not path or not path.exists(): + raise ValueError(f"Target file does not exist: {path}") - bbts_path = segment_path.with_suffix(".bbts") - ents_path = segment_path.with_suffix(".ents") + if not self._tool_path or not self._tool_path.exists(): + return + + if not self.key or not self.kid: + raise MonaLisa.Exceptions.KeyExtractionFailed("Missing content key for decryption.") + + target_output = output_path or path + temp_enc = path.with_name(f"{path.name}.temp_enc") try: - if segment_path.exists(): - segment_path.replace(bbts_path) + if path.exists(): + path.replace(temp_enc) else: - raise MonaLisa.Exceptions.DecryptionFailed(f"Segment file does not exist: {segment_path}") + raise FileNotFoundError(f"Target file disappeared: {path}") - cmd = [str(worker_path), str(self._key), str(bbts_path), str(ents_path)] + cmd = [ + str(self._tool_path), + "-i", + str(temp_enc), + "-o", + str(target_output), + "-k", + self.key_pair, + ] startupinfo = None if sys.platform == "win32": startupinfo = subprocess.STARTUPINFO() startupinfo.dwFlags |= subprocess.STARTF_USESHOWWINDOW - worker_timeout_s = 60 process = subprocess.run( cmd, stdout=subprocess.PIPE, @@ -245,50 +262,15 @@ class MonaLisa: encoding="utf-8", errors="replace", startupinfo=startupinfo, - timeout=worker_timeout_s, + timeout=60, ) - if process.returncode != 0: - raise MonaLisa.Exceptions.DecryptionFailed( - f"ML-Worker failed for {segment_path.name}: {process.stderr}" - ) + if process.returncode != 0 or not target_output.exists(): + raise RuntimeError(f"MonaLisa decryption failed for {path.name}: {process.stderr or 'unknown error'}") - if not ents_path.exists(): - raise MonaLisa.Exceptions.DecryptionFailed( - f"Decrypted .ents file was not created for {segment_path.name}" - ) - - with open(ents_path, "rb") as f: - ents_data = f.read() - - crypto = AES.new(self._aes_key, AES.MODE_ECB) - decrypted_data = unpad(crypto.decrypt(ents_data), AES.block_size) - - with open(segment_path, "wb") as f: - f.write(decrypted_data) - - except MonaLisa.Exceptions.DecryptionFailed: - raise - except subprocess.TimeoutExpired as e: - log.error("ML-Worker timed out after %ss for %s", worker_timeout_s, segment_path.name) - raise MonaLisa.Exceptions.DecryptionFailed( - f"ML-Worker timed out after {worker_timeout_s}s for {segment_path.name}" - ) from e - except Exception as e: - raise MonaLisa.Exceptions.DecryptionFailed(f"Failed to decrypt segment {segment_path.name}: {e}") finally: - if ents_path.exists(): - os.remove(ents_path) - if bbts_path != segment_path and bbts_path.exists(): - os.remove(bbts_path) - - def decrypt(self, _path: Path) -> None: - """ - MonaLisa uses per-segment decryption during download, through the - on_segment_downloaded callback. When the caller calls this method, - unshackle has already decrypted the media and muxed it into a container. - - Args: - _path: Path to the file (ignored). - """ - pass + if temp_enc.exists(): + try: + temp_enc.unlink() + except Exception: + pass diff --git a/unshackle/core/service.py b/unshackle/core/service.py index c654026b..804ef57b 100644 --- a/unshackle/core/service.py +++ b/unshackle/core/service.py @@ -405,6 +405,14 @@ class Service(metaclass=ABCMeta): session.mount("http://", session.adapters["https://"]) return session + @staticmethod + def get_binaries() -> list[dict]: + """ + Declare custom binary dependencies required by this service. + :returns: List of dicts specifying name, candidates, desc, etc. + """ + return [] + def authenticate(self, cookies: Optional[CookieJar] = None, credential: Optional[Credential] = None) -> None: """ Authenticate the Service with Cookies and/or Credentials (Email/Username and Password). diff --git a/unshackle/core/services.py b/unshackle/core/services.py index 94b5665a..e2ba18f6 100644 --- a/unshackle/core/services.py +++ b/unshackle/core/services.py @@ -10,6 +10,7 @@ from typing import Any, Iterable import click +from unshackle.core import binaries from unshackle.core.config import config from unshackle.core.service import Service from unshackle.core.service_repo import DirtyServiceRepo, is_repo_spec, refresh_repo, resolve_service_repo @@ -102,7 +103,26 @@ def load_services(paths: list[Path]) -> tuple[dict[str, object], list[str]]: return modules, errors +def register_service_binaries(modules: dict[str, object]) -> None: + """Register custom binary dependencies declared by services via get_binaries().""" + + for tag, service_cls in modules.items(): + get_binaries_fn = getattr(service_cls, "get_binaries", None) + if callable(get_binaries_fn): + try: + for dep in get_binaries_fn(): + name = dep.get("name") + if not name: + continue + candidates = dep.get("candidates") or (name.lower(),) + desc = dep.get("desc", f"{tag} dependency") + binaries.register(name, *candidates, desc=desc) + except Exception as e: + log.warning(f"Failed to register custom binaries for service {tag}: {e}") + + MODULES, LOAD_ERRORS = load_services(SERVICES) +register_service_binaries(MODULES) ALIASES = {tag: getattr(module, "ALIASES", ()) for tag, module in MODULES.items()} @@ -143,6 +163,7 @@ def reload_services(tags: Iterable[str]) -> list[str]: except Exception as e: errors.append(str(e)) LOAD_ERRORS.extend(errors) + register_service_binaries(MODULES) PENDING.difference_update(wanted) return errors