diff --git a/unshackle/core/api/download_manager.py b/unshackle/core/api/download_manager.py index 7e2f0a4..a52769c 100644 --- a/unshackle/core/api/download_manager.py +++ b/unshackle/core/api/download_manager.py @@ -43,6 +43,7 @@ class DownloadJob: output_files: List[str] = field(default_factory=list) error_message: Optional[str] = None error_details: Optional[str] = None + decryption_keys: Optional[Dict[str, Any]] = None # Cancellation support cancel_event: threading.Event = field(default_factory=threading.Event) @@ -67,6 +68,7 @@ class DownloadJob: "output_files": self.output_files, "error_message": self.error_message, "error_details": self.error_details, + "decryption_keys": self.decryption_keys, } ) @@ -80,8 +82,14 @@ def _perform_download( params: Dict[str, Any], cancel_event: Optional[threading.Event] = None, progress_callback: Optional[Callable[[Dict[str, Any]], None]] = None, -) -> List[str]: - """Execute the synchronous download logic for a job.""" +) -> tuple[List[str], Optional[Dict[str, Any]]]: + """Execute the synchronous download logic for a job. + + Returns: + Tuple of (output_files, decryption_keys) + - output_files: List of downloaded file paths + - decryption_keys: Dict of keys when skip_dl=True, None otherwise + """ def _check_cancel(stage: str): if cancel_event and cancel_event.is_set(): @@ -97,12 +105,20 @@ def _perform_download( import yaml from unshackle.commands.dl import dl from unshackle.core.config import config + from unshackle.core.constants import DOWNLOAD_LICENCE_ONLY from unshackle.core.services import Services from unshackle.core.utils.click_types import ContextData from unshackle.core.utils.collections import merge_dict log.info(f"Starting sync download for job {job_id}") + skip_dl = params.get("skip_dl", False) + collected_keys = None + + if skip_dl: + DOWNLOAD_LICENCE_ONLY.set() + log.info(f"Skip download mode enabled for job {job_id}, keys will be collected from DRM objects") + # Load service configuration service_config_path = Services.get_path(service) / config.filenames.config if service_config_path.exists(): @@ -179,6 +195,13 @@ def _perform_download( original_download_dir = config.directories.downloads _check_cancel("before download execution") + wanted_param = params.get("wanted", []) + if wanted_param and isinstance(wanted_param, str): + from unshackle.core.utils.click_types import SEASON_RANGE + + wanted_param = SEASON_RANGE.parse_tokens(wanted_param) + elif not wanted_param: + wanted_param = [] stdout_capture = StringIO() stderr_capture = StringIO() @@ -208,6 +231,28 @@ def _perform_download( dl_instance.result = result_with_progress + processed_titles = [] + collected_keys = {} + download_error = None + original_track_download = None + + if skip_dl: + from unshackle.core.tracks.track import Track + + original_track_download = Track.download + + def capturing_track_download(self, *args, **kwargs): + prepare_drm = kwargs.get("prepare_drm") + + if prepare_drm and hasattr(prepare_drm, "keywords"): + title = prepare_drm.keywords.get("title") + if title and title not in processed_titles: + processed_titles.append(title) + + return original_track_download(self, *args, **kwargs) + + Track.download = capturing_track_download + try: with redirect_stdout(stdout_capture), redirect_stderr(stderr_capture): dl_instance.result( @@ -220,7 +265,7 @@ def _perform_download( range_=params.get("range", []), channels=params.get("channels"), no_atmos=params.get("no_atmos", False), - wanted=params.get("wanted", []), + wanted=wanted_param, lang=params.get("lang", ["orig"]), v_lang=params.get("v_lang", []), a_lang=params.get("a_lang", []), @@ -256,7 +301,9 @@ def _perform_download( log.error(f"Download exited with code {exc.code}") log.error(f"Stdout: {stdout_str}") log.error(f"Stderr: {stderr_str}") - raise Exception(f"Download failed with exit code {exc.code}") + download_error = Exception(f"Download failed with exit code {exc.code}") + if not skip_dl: + raise download_error except Exception as exc: # noqa: BLE001 - propagate to caller stdout_str = stdout_capture.getvalue() @@ -264,11 +311,85 @@ def _perform_download( log.error(f"Download execution failed: {exc}") log.error(f"Stdout: {stdout_str}") log.error(f"Stderr: {stderr_str}") - raise + download_error = exc + if not skip_dl: + raise - log.info(f"Download completed for job {job_id}, files in {original_download_dir}") + finally: + # Clear the DOWNLOAD_LICENCE_ONLY event after download completes + if skip_dl: + # Restore original Track.download method + if original_track_download is not None: + from unshackle.core.tracks.track import Track - return [] + Track.download = original_track_download + + DOWNLOAD_LICENCE_ONLY.clear() + log.info(f"Cleared skip download mode for job {job_id}") + + # Extract keys directly from processed titles' tracks + log.debug(f"Processing {len(processed_titles)} captured titles for key extraction") + + if processed_titles: + try: + for title in processed_titles: + title_name = str(title) + collected_keys[title_name] = {} + + # Extract keys from all tracks that have DRM + for track in title.tracks: + if not hasattr(track, "drm") or not track.drm: + continue + + # track.drm can be a single DRM object or a list of DRM objects + drm_objects = track.drm if isinstance(track.drm, list) else [track.drm] + track_name = str(track) + track_keys = {} + + # Extract keys from each DRM object + for drm in drm_objects: + if not drm or not hasattr(drm, "content_keys"): + continue + + if drm.content_keys: + # Convert UUID keys to hex strings for JSON serialization + for kid, key in drm.content_keys.items(): + kid_hex = kid.hex if hasattr(kid, "hex") else str(kid) + track_keys[kid_hex] = key + + if track_keys: + collected_keys[title_name][track_name] = track_keys + log.debug(f"Extracted {len(track_keys)} key(s) from {track_name}") + + # Remove title if no keys were collected + if not collected_keys[title_name]: + del collected_keys[title_name] + + if collected_keys: + log.debug(f"Collected keys for {len(collected_keys)} title(s) from DRM objects") + else: + log.warning("No keys found in processed titles' DRM objects") + + except Exception as e: + log.error(f"Failed to extract keys from processed titles: {e}") + import traceback + + log.error(traceback.format_exc()) + else: + log.warning("No titles were captured during processing") + + if skip_dl: + if collected_keys: + log.info(f"Key extraction completed for job {job_id} - {len(collected_keys)} title(s)") + else: + log.warning(f"Key extraction completed for job {job_id} but no keys were found") + + if download_error and not collected_keys: + raise download_error + else: + log.info(f"Download completed for job {job_id}, files in {original_download_dir}") + + return [], collected_keys class DownloadQueueManager: @@ -472,11 +593,15 @@ class DownloadQueueManager: log.info(f"Executing download for job {job.job_id}") try: - output_files = await self._run_download_async(job) + output_files, decryption_keys = await self._run_download_async(job) job.status = JobStatus.COMPLETED job.output_files = output_files + job.decryption_keys = decryption_keys job.progress = 100.0 - log.info(f"Download completed for job {job.job_id}: {len(output_files)} files") + if decryption_keys: + log.info(f"Download completed for job {job.job_id}: retrieved keys for {len(decryption_keys)} title(s)") + else: + log.info(f"Download completed for job {job.job_id}: {len(output_files)} files") except Exception as e: job.status = JobStatus.FAILED job.error_message = str(e) @@ -484,8 +609,12 @@ class DownloadQueueManager: log.error(f"Download failed for job {job.job_id}: {e}") raise - async def _run_download_async(self, job: DownloadJob) -> List[str]: - """Invoke a worker subprocess to execute the download.""" + async def _run_download_async(self, job: DownloadJob) -> tuple[List[str], Optional[Dict[str, Any]]]: + """Invoke a worker subprocess to execute the download. + + Returns: + Tuple of (output_files, decryption_keys) + """ payload = { "job_id": job.job_id, @@ -580,7 +709,9 @@ class DownloadQueueManager: message = result_data.get("message") if result_data else "worker did not report success" raise Exception(f"Worker failure: {message}") - return result_data.get("output_files", []) + output_files = result_data.get("output_files", []) + decryption_keys = result_data.get("decryption_keys") + return output_files, decryption_keys finally: if not communicate_task.done(): @@ -597,7 +728,7 @@ class DownloadQueueManager: except OSError: pass - def _execute_download_sync(self, job: DownloadJob) -> List[str]: + def _execute_download_sync(self, job: DownloadJob) -> tuple[List[str], Optional[Dict[str, Any]]]: """Execute download synchronously using existing dl.py logic.""" return _perform_download(job.job_id, job.service, job.title_id, job.parameters.copy(), job.cancel_event) diff --git a/unshackle/core/api/download_worker.py b/unshackle/core/api/download_worker.py index 08810d4..af4bf23 100644 --- a/unshackle/core/api/download_worker.py +++ b/unshackle/core/api/download_worker.py @@ -59,11 +59,13 @@ def main(argv: list[str]) -> int: except Exception as e: log.error(f"Failed to write progress update: {e}") - output_files = _perform_download( + output_files, decryption_keys = _perform_download( job_id, service, title_id, params, cancel_event=None, progress_callback=progress_callback ) result = {"status": "success", "output_files": output_files} + if decryption_keys: + result["decryption_keys"] = decryption_keys except Exception as exc: # noqa: BLE001 - capture for parent process exit_code = 1 diff --git a/unshackle/core/api/handlers.py b/unshackle/core/api/handlers.py index 60261a6..5cd4bed 100644 --- a/unshackle/core/api/handlers.py +++ b/unshackle/core/api/handlers.py @@ -559,8 +559,58 @@ async def list_tracks_handler(data: Dict[str, Any]) -> web.Response: return web.json_response({"status": "error", "message": str(e)}, status=500) +async def keys_handler(data: Dict[str, Any]) -> web.Response: + """Handle keys request - retrieve decryption keys without downloading. + + Similar to list-tracks but performs license acquisition and returns keys. + """ + service_tag = data.get("service") + title_id = data.get("title_id") + data.get("profile") + + if not service_tag: + return web.json_response({"status": "error", "message": "Missing required parameter: service"}, status=400) + + if not title_id: + return web.json_response({"status": "error", "message": "Missing required parameter: title_id"}, status=400) + + normalized_service = validate_service(service_tag) + if not normalized_service: + return web.json_response( + {"status": "error", "message": f"Invalid or unavailable service: {service_tag}"}, status=400 + ) + + try: + from unshackle.core.api.download_manager import _perform_download + import uuid + + temp_job_id = str(uuid.uuid4()) + + params = {k: v for k, v in data.items() if k not in ["service", "title_id"]} + params["skip_dl"] = True + + output_files, decryption_keys = _perform_download( + temp_job_id, normalized_service, title_id, params, cancel_event=None, progress_callback=None + ) + + if decryption_keys: + return web.json_response({"keys": decryption_keys}) + else: + return web.json_response( + {"status": "error", "message": "No decryption keys found for this title"}, status=404 + ) + + except Exception as e: + log.exception("Error retrieving keys") + return web.json_response({"status": "error", "message": str(e)}, status=500) + + async def download_handler(data: Dict[str, Any]) -> web.Response: - """Handle download request - create and queue a download job.""" + """Handle download request - create and queue a download job. + + Supports skip_dl parameter to retrieve decryption keys without downloading tracks. + When skip_dl=True, the job will obtain licenses and keys but skip actual track downloads. + """ from unshackle.core.api.download_manager import get_download_manager service_tag = data.get("service") diff --git a/unshackle/core/api/routes.py b/unshackle/core/api/routes.py index 5445c87..769d4c3 100644 --- a/unshackle/core/api/routes.py +++ b/unshackle/core/api/routes.py @@ -4,8 +4,15 @@ from aiohttp import web from aiohttp_swagger3 import SwaggerDocs, SwaggerInfo, SwaggerUiSettings from unshackle.core import __version__ -from unshackle.core.api.handlers import (cancel_download_job_handler, download_handler, get_download_job_handler, - list_download_jobs_handler, list_titles_handler, list_tracks_handler) +from unshackle.core.api.handlers import ( + cancel_download_job_handler, + download_handler, + get_download_job_handler, + keys_handler, + list_download_jobs_handler, + list_titles_handler, + list_tracks_handler, +) from unshackle.core.services import Services from unshackle.core.update_checker import UpdateChecker @@ -226,6 +233,73 @@ async def list_tracks(request: web.Request) -> web.Response: return await list_tracks_handler(data) +async def keys(request: web.Request) -> web.Response: + """ + Retrieve decryption keys for a title without downloading. + --- + summary: Get decryption keys + description: Retrieve decryption keys for a title without downloading content + requestBody: + required: true + content: + application/json: + schema: + type: object + required: + - service + - title_id + properties: + service: + type: string + description: Service tag + title_id: + type: string + description: Title identifier + wanted: + type: string + description: Specific episode/season (optional) + proxy: + type: string + description: Proxy configuration (optional) + quality: + type: array + items: + type: integer + description: Quality levels (optional) + lang: + type: array + items: + type: string + description: Audio languages (optional) + s_lang: + type: array + items: + type: string + description: Subtitle languages (optional) + responses: + '200': + description: Decryption keys + content: + application/json: + schema: + type: object + properties: + keys: + type: object + description: Map of KID to decryption key (hex format) + '400': + description: Invalid request + '500': + description: Server error + """ + try: + data = await request.json() + except Exception: + return web.json_response({"status": "error", "message": "Invalid JSON request body"}, status=400) + + return await keys_handler(data) + + async def download(request: web.Request) -> web.Response: """ Download content based on provided parameters. @@ -355,6 +429,7 @@ def setup_routes(app: web.Application) -> None: app.router.add_get("/api/services", services) app.router.add_post("/api/list-titles", list_titles) app.router.add_post("/api/list-tracks", list_tracks) + app.router.add_post("/api/keys", keys) app.router.add_post("/api/download", download) app.router.add_get("/api/download/jobs", download_jobs) app.router.add_get("/api/download/jobs/{job_id}", download_job_detail) @@ -380,6 +455,7 @@ def setup_swagger(app: web.Application) -> None: web.get("/api/services", services), web.post("/api/list-titles", list_titles), web.post("/api/list-tracks", list_tracks), + web.post("/api/keys", keys), web.post("/api/download", download), web.get("/api/download/jobs", download_jobs), web.get("/api/download/jobs/{job_id}", download_job_detail),