mirror of
https://github.com/unshackle-dl/unshackle.git
synced 2026-03-12 09:29:02 +00:00
feat(api): add /api/keys endpoint for synchronous decryption key retrieval
Add a new /api/keys endpoint that retrieves DRM decryption keys without
downloading content, similar to CLI --skip-dl
Key changes:
- Add keys_handler in handlers.py for synchronous key extraction
- Add /api/keys POST route with full Swagger documentation
- Implement title/track capture mechanism via Track.download patching
- Handle both single DRM objects and lists of DRM objects per track
- Parse wanted parameter (S01E01 -> 1x1 format) for episode filtering
- Return partial results if some episodes fail (e.g., not yet available)
Response format:
{
"keys": {
"Title Name": {
"Track Name": {
"kid_hex": "key_hex"
}
}
}
}
This commit is contained in:
@@ -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)
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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),
|
||||
|
||||
Reference in New Issue
Block a user