diff --git a/unshackle/commands/dl.py b/unshackle/commands/dl.py index 52f7d66b..dc69c17c 100644 --- a/unshackle/commands/dl.py +++ b/unshackle/commands/dl.py @@ -1,6 +1,7 @@ from __future__ import annotations import html +import json import logging import math import os @@ -22,7 +23,6 @@ from typing import Any, Callable, Optional from uuid import UUID import click -import jsonpickle import yaml from langcodes import Language from pymediainfo import MediaInfo @@ -56,29 +56,14 @@ from unshackle.core.tracks import Audio, Subtitle, Tracks, Video from unshackle.core.tracks.attachment import Attachment from unshackle.core.tracks.dv_fixup import apply_dv_fixup from unshackle.core.tracks.hybrid import Hybrid -from unshackle.core.utilities import ( - find_font_with_fallbacks, - find_missing_langs, - get_debug_logger, - get_system_fonts, - init_debug_logger, - is_close_match, - suggest_font_packages, - time_elapsed_since, -) +from unshackle.core.utilities import (find_font_with_fallbacks, find_missing_langs, get_debug_logger, + get_system_fonts, init_debug_logger, is_close_match, suggest_font_packages, + time_elapsed_since) from unshackle.core.utils import tags from unshackle.core.utils.bitrate import apply_real_bitrates -from unshackle.core.utils.click_types import ( - AUDIO_CODEC_LIST, - LANGUAGE_RANGE, - QUALITY_LIST, - SEASON_RANGE, - SLOW_DELAY_RANGE, - ContextData, - MultipleChoice, - MultipleVideoCodecChoice, - SubtitleCodecChoice, -) +from unshackle.core.utils.click_types import (AUDIO_CODEC_LIST, LANGUAGE_RANGE, QUALITY_LIST, SEASON_RANGE, + SLOW_DELAY_RANGE, ContextData, MultipleChoice, MultipleVideoCodecChoice, + SubtitleCodecChoice) from unshackle.core.utils.collections import merge_dict from unshackle.core.utils.selector import select_multiple from unshackle.core.utils.subprocess import ffprobe @@ -523,6 +508,14 @@ class dl: default=False, help="Export track info and decryption keys to a JSON file in the exports directory.", ) + @click.option( + "--import", + "import_file", + type=str, + default=None, + hidden=True, + help="Internal: path to an export JSON to reconstruct a download from (used by 'unshackle import').", + ) @click.option( "--cdm-only/--vaults-only", is_flag=True, @@ -658,6 +651,7 @@ class dl: ) self.profile = profile + self.proxy_requested = bool(proxy) self.tmdb_id = tmdb_id self.imdb_id = imdb_id self.enrich = enrich @@ -1100,6 +1094,7 @@ class dl: if export: config.directories.exports.mkdir(parents=True, exist_ok=True) export_path = config.directories.exports / f"export_{self.service}_{int(time.time())}.json" + self.export_service = service else: export_path = None @@ -2761,52 +2756,80 @@ class dl: console.print(Padding(f"Processed all titles in [progress.elapsed]{dl_time}", (0, 5, 1, 5))) - def _write_export(self, export: Path, title: Title_T, track: AnyTrack, drm: Any) -> None: - """Write decryption keys and track info to the export JSON file.""" + @staticmethod + def title_to_meta(title: Title_T) -> dict[str, Any]: + """Capture the title fields needed to rebuild a Title for output naming on import.""" + from unshackle.core.titles import Episode, Movie + + meta: dict[str, Any] = { + "id": str(title.id), + "language": str(title.language) if getattr(title, "language", None) else None, + } + if isinstance(title, Episode): + meta.update( + type="episode", + series_title=title.title, + season=title.season, + number=title.number, + name=title.name, + year=title.year, + ) + elif isinstance(title, Movie): + meta.update(type="movie", name=title.name, year=title.year) + else: + meta.update(type="movie", name=str(title)) + return meta + + def write_export(self, export: Path, title: Title_T, track: AnyTrack, drm: Any) -> None: + """Write a shareable v2 export usable by ``unshackle import``. + + Carries no session/cookies/dl-flags. Region (country code) is stored only when the + export used ``--proxy``, as an import geofence. Each track records only the licensed + DRM system; content keys live once under the track's ``keys``. + """ with self.EXPORT_LOCK: - keys = {} + doc: dict[str, Any] = {} if export.is_file(): - keys = jsonpickle.loads(export.read_text(encoding="utf8")) or {} - if str(title) not in keys: - keys[str(title)] = {} + doc = json.loads(export.read_text(encoding="utf8")) or {} - title_data = keys[str(title)] + doc.setdefault("version", 2) + doc.setdefault("service", self.service) + if "region" not in doc and getattr(self, "proxy_requested", False): + region = getattr(getattr(self, "export_service", None), "current_region", None) + if region: + doc["region"] = region - if title.tracks.manifest_url and "manifest" not in title_data: - title_data["manifest"] = title.tracks.manifest_url + titles = doc.setdefault("titles", {}) + tinfo = titles.setdefault(str(title.id), {}) + tinfo.setdefault("meta", self.title_to_meta(title)) - if isinstance(track, Video): - section = "video" - elif isinstance(track, Audio): - section = "audio" - else: - section = "other" + if title.tracks.manifest_url: + tinfo.setdefault("manifest_url", title.tracks.manifest_url) - if section not in title_data: - title_data[section] = {} - if str(track) not in title_data[section]: - title_data[section][str(track)] = {} + all_tracks = [*title.tracks.videos, *title.tracks.audio, *title.tracks.subtitles] + if "manifest_type" not in tinfo: + tinfo["manifest_type"] = next( + (t.descriptor.name for t in all_tracks if t.descriptor != Video.Descriptor.URL), None + ) - track_data = title_data[section][str(track)] - track_data["url"] = track.url - track_data["descriptor"] = track.descriptor.name + tracks_map = tinfo.setdefault("tracks", {}) + if not tracks_map: + for t in all_tracks: + tracks_map[str(t.id)] = t.to_dict() - if "keys" not in track_data: - track_data["keys"] = {} + track_data = tracks_map.setdefault(str(track.id), track.to_dict()) + if hasattr(drm, "to_dict"): + track_data["drm"] = [drm.to_dict()] + keys = track_data.setdefault("keys", {}) for kid, key in drm.content_keys.items(): - track_data["keys"][kid.hex] = key + keys[kid.hex] = key - if "subtitles" not in title_data: - subs = {} - for sub in title.tracks.subtitles: - subs[str(sub)] = {"url": sub.url} - if subs: - title_data["subtitles"] = subs + if "chapters" not in tinfo: + tinfo["chapters"] = [ + {"timestamp": chapter.timestamp, "name": chapter.name} for chapter in (title.tracks.chapters or []) + ] - section_order = ["manifest", "video", "audio", "subtitles", "other"] - keys[str(title)] = {k: title_data[k] for k in section_order if k in title_data} - - export.write_text(jsonpickle.dumps(keys, indent=4), encoding="utf8") + export.write_text(json.dumps(doc, indent=4, ensure_ascii=False), encoding="utf8") def prepare_drm( self, @@ -3088,7 +3111,7 @@ class dl: table.add_row(cek_tree) if export: - self._write_export(export, title, track, drm) + self.write_export(export, title, track, drm) elif isinstance(drm, PlayReady): if self.debug_logger: @@ -3259,7 +3282,7 @@ class dl: table.add_row(cek_tree) if export: - self._write_export(export, title, track, drm) + self.write_export(export, title, track, drm) elif isinstance(drm, MonaLisa): with self.DRM_TABLE_LOCK: diff --git a/unshackle/commands/import.py b/unshackle/commands/import.py new file mode 100644 index 00000000..79ec5bec --- /dev/null +++ b/unshackle/commands/import.py @@ -0,0 +1,54 @@ +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any + +import click + +from unshackle.commands.dl import dl +from unshackle.core.constants import context_settings + + +class ImportCommand: + @staticmethod + @click.command( + name="import", + short_help="Reconstruct a download (download, decrypt, mux) from an --export JSON file.", + context_settings={**context_settings, "ignore_unknown_options": True}, + ) + @click.argument("export_file", type=Path) + @click.argument("dl_args", nargs=-1, type=click.UNPROCESSED) + @click.pass_context + def cli(ctx: click.Context, export_file: Path, dl_args: tuple[str, ...]) -> None: + """ + Reconstruct an exported download without re-contacting the service. + + Re-fetches the manifest, injects the stored keys, then downloads/decrypts/muxes as a + normal `dl` run. Any `dl` options after the file are forwarded verbatim: + + unshackle import export.json -r HDR10 --proxy US + """ + if not export_file.is_file(): + raise click.ClickException(f"Export file not found: {export_file}") + + try: + data: dict[str, Any] = json.loads(export_file.read_text(encoding="utf8")) + except json.JSONDecodeError as e: + raise click.ClickException(f"Export file is not valid JSON: {e}") + + if data.get("version") != 2: + raise click.ClickException( + f"Unsupported export version {data.get('version')!r}. " + "Re-create the export with a current build of unshackle." + ) + + service_tag = data.get("service") + if not service_tag: + raise click.ClickException("Export file is missing the 'service' tag.") + + args = [*dl_args, "--import", str(export_file), service_tag] + dl.cli.main(args=args, prog_name="unshackle dl", standalone_mode=False) + + +globals()["import"] = ImportCommand diff --git a/unshackle/core/drm/__init__.py b/unshackle/core/drm/__init__.py index d94e9129..bfe126d2 100644 --- a/unshackle/core/drm/__init__.py +++ b/unshackle/core/drm/__init__.py @@ -1,4 +1,6 @@ -from typing import Union +import base64 +from typing import Any, Union +from uuid import UUID from unshackle.core.drm.clearkey import ClearKey from unshackle.core.drm.monalisa import MonaLisa @@ -8,4 +10,35 @@ from unshackle.core.drm.widevine import Widevine DRM_T = Union[ClearKey, Widevine, PlayReady, MonaLisa] -__all__ = ("ClearKey", "Widevine", "PlayReady", "MonaLisa", "DRM_T") +def drm_from_dict(data: dict[str, Any]) -> Union[Widevine, PlayReady]: + """Reconstruct a Widevine/PlayReady DRM instance from its ``to_dict()`` form. + + Rebuilds the PSSH from the stored base64 and re-injects any saved content keys + so the resulting object can decrypt without contacting a license server. + """ + system = data.get("system") + pssh_b64 = data.get("pssh_b64") + kids = data.get("kids") or [] + content_keys = data.get("content_keys") or {} + + if not pssh_b64: + raise ValueError("Cannot reconstruct DRM without a stored PSSH.") + + if system == "PlayReady": + from pyplayready.system.pssh import PSSH as PlayReadyPSSH + + drm: Union[Widevine, PlayReady] = PlayReady(pssh=PlayReadyPSSH(base64.b64decode(pssh_b64)), pssh_b64=pssh_b64) + elif system == "Widevine": + from pywidevine.pssh import PSSH as WidevinePSSH + + drm = Widevine(pssh=WidevinePSSH(pssh_b64), kid=kids[0] if kids else None) + else: + raise ValueError(f"Unsupported DRM system for reconstruction: {system!r}") + + for kid_hex, key in content_keys.items(): + drm.content_keys[UUID(hex=kid_hex)] = key + + return drm + + +__all__ = ("ClearKey", "Widevine", "PlayReady", "MonaLisa", "DRM_T", "drm_from_dict") diff --git a/unshackle/core/drm/playready.py b/unshackle/core/drm/playready.py index 253e4a57..09d0756d 100644 --- a/unshackle/core/drm/playready.py +++ b/unshackle/core/drm/playready.py @@ -247,6 +247,17 @@ class PlayReady: def kid(self) -> Optional[UUID]: return next(iter(self.kids), None) + def to_dict(self) -> dict[str, Any]: + """Serialise this DRM instance for export/import (PSSH + KIDs). + + Content keys are stored once at the export's track level, not duplicated here. + """ + return { + "system": "PlayReady", + "pssh_b64": self.pssh_b64, + "kids": [kid.hex for kid in self.kids], + } + @property def kids(self) -> list[UUID]: return self._kids diff --git a/unshackle/core/drm/widevine.py b/unshackle/core/drm/widevine.py index 0bad0a67..250b05a5 100644 --- a/unshackle/core/drm/widevine.py +++ b/unshackle/core/drm/widevine.py @@ -180,6 +180,17 @@ class Widevine: return [self._kid] return [] + def to_dict(self) -> dict[str, Any]: + """Serialise this DRM instance for export/import (PSSH + KIDs). + + Content keys are stored once at the export's track level, not duplicated here. + """ + return { + "system": "Widevine", + "pssh_b64": self.pssh.dumps(), + "kids": [kid.hex for kid in self.kids], + } + def get_content_keys(self, cdm: WidevineCdm, certificate: Callable, licence: Callable) -> None: """ Create a CDM Session and obtain Content Keys for this DRM Instance. diff --git a/unshackle/core/import_service.py b/unshackle/core/import_service.py new file mode 100644 index 00000000..d62ff054 --- /dev/null +++ b/unshackle/core/import_service.py @@ -0,0 +1,284 @@ +from __future__ import annotations + +import json +import logging +from http.cookiejar import CookieJar +from pathlib import Path +from typing import Any, Optional, Union +from uuid import UUID + +import click +import requests + +from unshackle.core.config import config +from unshackle.core.constants import AnyTrack +from unshackle.core.credential import Credential +from unshackle.core.drm import drm_from_dict +from unshackle.core.manifests import DASH, HLS, ISM +from unshackle.core.remote_service import RemoteService, _build_title, _resolve_proxy +from unshackle.core.titles import Episode, Movies, Series, Title_T, Titles_T +from unshackle.core.tracks import Audio, Chapter, Chapters, Tracks, Video +from unshackle.core.tracks.track import Track + +log = logging.getLogger("import") + +PARSERS = {"DASH": DASH, "HLS": HLS, "ISM": ISM} + + +class ImportService: + """Reconstructs a download from an export JSON. + + Auth and licensing are skipped; tracks are rebuilt from the export and keys injected + directly. ``_server_cdm``/``_server_cdm_type`` keep their underscores: dl.py reads them + via getattr as the server-CDM contract that skips client licensing. + """ + + ALIASES: tuple[str, ...] = () + GEOFENCE: tuple[str, ...] = () + NO_SUBTITLES: bool = False + + def __init__(self, ctx: click.Context, service_tag: str, title: str, import_file: Optional[str]) -> None: + self.__class__.__name__ = service_tag + self.service_tag = service_tag + self.title_id = title + self.ctx = ctx + self.log = logging.getLogger(service_tag) + self.credential: Optional[Credential] = None + self.current_region: Optional[str] = None + self.title_cache = None + + if not import_file: + raise click.ClickException("No export file was provided to import from.") + export_path = Path(import_file) + if not export_path.is_file(): + raise click.ClickException(f"Export file not found: {export_path}") + + self.data: dict[str, Any] = json.loads(export_path.read_text(encoding="utf8")) + version = self.data.get("version") + if version != 2: + raise click.ClickException( + f"Unsupported export version {version!r}. Re-create the export with a current build." + ) + + self.titles_data: dict[str, Any] = self.data.get("titles", {}) + self.region: Optional[str] = self.data.get("region") + self.titles: Optional[Titles_T] = None + self.tracks_by_title: dict[str, Tracks] = {} + + self._server_cdm = True + self._server_cdm_type = "widevine" + + self.session = self.build_session(ctx, self.region) + + @staticmethod + def build_session(ctx: click.Context, region: Optional[str] = None) -> requests.Session: + """Session for re-fetching the manifest. + + Honours the importer's ``--proxy``; otherwise falls back to the export region as a + geofence. An explicit proxy that fails to resolve raises; a region fallback warns. + """ + session = requests.Session() + session.headers.update(config.headers) + + params = ctx.parent.params if ctx.parent else {} + if params.get("no_proxy"): + return session + + explicit = params.get("proxy") + proxy_query = explicit or region + if not proxy_query: + return session + + try: + proxy = _resolve_proxy(proxy_query) + except Exception as e: + if explicit: + raise click.ClickException(f"Failed to resolve proxy '{proxy_query}': {e}") + log.warning(f"Could not auto-select a proxy for export region '{region}': {e}. Continuing without proxy.") + proxy = None + + if proxy: + session.proxies.update({"all": proxy}) + if not explicit: + log.info(f"No --proxy given; using export region '{region}' via your proxy provider.") + return session + + @property + def title(self) -> str: + return self.title_id + + def authenticate(self, cookies: Optional[CookieJar] = None, credential: Optional[Credential] = None) -> None: + self.credential = credential + + def get_titles(self) -> Titles_T: + if self.titles is not None: + return self.titles + titles_list = [ + _build_title(entry.get("meta", {}), self.service_tag, fallback_id=title_id) + for title_id, entry in self.titles_data.items() + ] + self.titles = ( + Series(titles_list) if titles_list and isinstance(titles_list[0], Episode) else Movies(titles_list) + ) + return self.titles + + def get_titles_cached(self, title_id: Optional[str] = None) -> Titles_T: + return self.get_titles() + + def get_tracks(self, title: Title_T) -> Tracks: + """Reconstruct the title's tracks from the export. + + DASH/ISM: re-fetch and re-parse the manifest and return the full ladder (the importer + picks quality with normal dl flags; keys are injected by KID later). HLS/URL: rebuild + from the stored per-track dicts, since the variant is re-fetched from track.url at + download time and ATV-style master playlists carry unstable per-fetch tokens. + """ + title_id = str(title.id) + if title_id in self.tracks_by_title: + return self.tracks_by_title[title_id] + + entry = self.titles_data.get(title_id, {}) + tracks_map: dict[str, Any] = entry.get("tracks") or {} + manifest_url = entry.get("manifest_url") + manifest_type = entry.get("manifest_type") + + tracks = Tracks() + tracks.manifest_url = manifest_url + + parser = PARSERS.get(manifest_type or "") + if manifest_url and parser is not None and manifest_type in ("DASH", "ISM"): + try: + parsed = parser.from_url(url=manifest_url, session=self.session).to_tracks(language=title.language) + except Exception as e: + raise click.ClickException( + f"Failed to re-fetch/parse the {manifest_type} manifest for '{title}'. " + f"The manifest URL may have expired since export. ({e})" + ) + for track in parsed: + tracks.add(track) + else: + for track_dict in tracks_map.values(): + track = Track.from_dict(track_dict) + drm = self.rebuild_drm(track_dict) + if drm: + track.drm = drm + tracks.add(track) + + self.tracks_by_title[title_id] = tracks + return tracks + + def key_pool(self) -> dict[UUID, str]: + """All exported KID:KEY pairs across every title, as {UUID: key_hex}.""" + pool: dict[UUID, str] = {} + for entry in self.titles_data.values(): + for track_dict in (entry.get("tracks") or {}).values(): + for kid_hex, key in (track_dict.get("keys") or {}).items(): + pool[UUID(hex=kid_hex)] = key + return pool + + def rebuild_drm(self, track_dict: dict[str, Any]) -> Optional[list[Any]]: + """Rebuild a DRM object (from stored PSSH, falling back to a stub) with the exported keys.""" + keys = track_dict.get("keys") or {} + drm_dicts = track_dict.get("drm") or [] + if not drm_dicts and not keys: + return None + + drm_obj = None + if drm_dicts: + try: + drm_obj = drm_from_dict(drm_dicts[0]) + except Exception as e: + self.log.debug(f"Falling back to DRM stub (PSSH rebuild failed: {e})") + + if drm_obj is None and keys: + drm_type = (drm_dicts[0].get("system", "Widevine").lower()) if drm_dicts else "widevine" + drm_obj = RemoteService._create_drm_stub(drm_type, list(keys.keys())) + + if drm_obj is None: + return None + + for kid_hex, key in keys.items(): + drm_obj.content_keys[UUID(hex=kid_hex)] = key + return [drm_obj] + + def resolve_server_keys(self, title: Title_T) -> None: + """Inject exported keys into the selected encrypted tracks by KID (no network). + + Called by dl.py after selection. Only encrypted video/audio are touched; encrypted + DASH tracks (no DRM at parse time) get a stub holding the keys, which + DASH.download_track preserves. decrypt() applies the key whose KID matches the media. + """ + pool = self.key_pool() + if not pool: + return + + system = self.exported_drm_system() + kid_hexes = [kid.hex for kid in pool] + + for track in title.tracks: + if not isinstance(track, (Video, Audio)) or not self.track_is_encrypted(track): + continue + drm_obj = track.drm[0] if track.drm else RemoteService._create_drm_stub(system, kid_hexes) + for kid, key in pool.items(): + drm_obj.content_keys[kid] = key + track.drm = [drm_obj] + self._server_cdm_type = drm_obj.__class__.__name__.lower() + + @staticmethod + def track_is_encrypted(track: Any) -> bool: + """True if the track carries DRM or its DASH manifest declares ContentProtection.""" + if track.drm: + return True + dash = track.data.get("dash") if getattr(track, "data", None) else None + if dash: + for element in (dash.get("representation"), dash.get("adaptation_set")): + if element is not None and element.findall("ContentProtection"): + return True + return False + + def exported_drm_system(self) -> str: + """The DRM system the exporter licensed (e.g. 'playready'), defaulting to widevine.""" + for entry in self.titles_data.values(): + for track_dict in (entry.get("tracks") or {}).values(): + for drm_dict in track_dict.get("drm") or []: + if drm_dict.get("system"): + return drm_dict["system"].lower() + return "widevine" + + def get_chapters(self, title: Title_T) -> Chapters: + entry = self.titles_data.get(str(title.id), {}) + return Chapters( + [Chapter(ch["timestamp"], ch.get("name")) for ch in (entry.get("chapters") or []) if ch.get("timestamp")] + ) + + def get_widevine_service_certificate(self, **_: Any) -> Optional[str]: + return None + + def get_widevine_license(self, *, challenge: bytes, title: Title_T, track: AnyTrack) -> Optional[Union[bytes, str]]: + raise RuntimeError("ImportService should not request a license; keys come from the export.") + + def get_playready_license( + self, *, challenge: bytes, title: Title_T, track: AnyTrack + ) -> Optional[Union[bytes, str]]: + raise RuntimeError("ImportService should not request a license; keys come from the export.") + + def on_segment_downloaded(self, track: AnyTrack, segment: Any) -> None: + pass + + def on_track_downloaded(self, track: AnyTrack) -> None: + pass + + def on_track_decrypted(self, track: AnyTrack, drm: Any, segment: Any = None) -> None: + pass + + def on_track_repacked(self, track: AnyTrack) -> None: + pass + + def on_track_multiplex(self, track: AnyTrack) -> None: + pass + + def close(self) -> None: + try: + self.session.close() + except Exception: + pass diff --git a/unshackle/core/services.py b/unshackle/core/services.py index c013d0da..a26d1799 100644 --- a/unshackle/core/services.py +++ b/unshackle/core/services.py @@ -70,6 +70,10 @@ class Services(click.MultiCommand): """Load the Service and return the Click CLI method.""" tag = Services.get_tag(name) + import_file = ctx.params.get("import_file") or (ctx.parent and ctx.parent.params.get("import_file")) + if import_file: + return Services._make_import_command(tag, ctx) + remote = ctx.params.get("remote") or (ctx.parent and ctx.parent.params.get("remote")) if remote: return Services._make_remote_command(tag, ctx) @@ -141,6 +145,24 @@ class Services(click.MultiCommand): return remote_cli + @staticmethod + def _make_import_command(tag: str, ctx: click.Context) -> click.Command: + """Create a synthetic command that yields an ImportService from an export JSON. + + Mirrors how remote services are wired so dl.py's result() runs unchanged. + """ + + @click.command(name=tag, short_help="Reconstruct a download from an export JSON.") + @click.argument("title", type=str, required=False, default="") + @click.pass_context + def import_cli(ctx: click.Context, title: str, **kwargs: object) -> object: + from unshackle.core.import_service import ImportService + + import_file = ctx.params.get("import_file") or (ctx.parent and ctx.parent.params.get("import_file")) + return ImportService(ctx, tag, title, import_file) + + return import_cli + @staticmethod def _fetch_remote_service_info(tag: str, ctx: click.Context) -> dict | None: """Fetch service info for a specific service from the remote server.""" diff --git a/unshackle/core/tracks/audio.py b/unshackle/core/tracks/audio.py index b4c69cd3..988136b2 100644 --- a/unshackle/core/tracks/audio.py +++ b/unshackle/core/tracks/audio.py @@ -126,6 +126,31 @@ class Audio(Track): self.joc = joc self.descriptive = bool(descriptive) + def to_dict(self) -> dict[str, Any]: + data = super().to_dict() + data.update( + { + "codec": self.codec.name if self.codec else None, + "bitrate": self.bitrate, + "channels": self.channels, + "joc": self.joc, + "descriptive": self.descriptive, + } + ) + return data + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> Audio: + kwargs = Track.base_kwargs_from_dict(data) + return cls( + **kwargs, + codec=Audio.Codec[data["codec"]] if data.get("codec") else None, + bitrate=data.get("bitrate"), + channels=data.get("channels"), + joc=data.get("joc"), + descriptive=data.get("descriptive", False), + ) + @property def atmos(self) -> bool: """Return True if the audio track contains Dolby Atmos.""" diff --git a/unshackle/core/tracks/subtitle.py b/unshackle/core/tracks/subtitle.py index 4ed84b40..847056ac 100644 --- a/unshackle/core/tracks/subtitle.py +++ b/unshackle/core/tracks/subtitle.py @@ -195,6 +195,29 @@ class Subtitle(Track): # Called after Track has been converted to another format self.OnConverted: Optional[Callable[[Subtitle.Codec], None]] = None + def to_dict(self) -> dict[str, Any]: + data = super().to_dict() + data.update( + { + "codec": self.codec.name if self.codec else None, + "cc": self.cc, + "sdh": self.sdh, + "forced": self.forced, + } + ) + return data + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> Subtitle: + kwargs = Track.base_kwargs_from_dict(data) + return cls( + **kwargs, + codec=Subtitle.Codec[data["codec"]] if data.get("codec") else None, + cc=data.get("cc", False), + sdh=data.get("sdh", False), + forced=data.get("forced", False), + ) + def __str__(self) -> str: return " | ".join( filter( diff --git a/unshackle/core/tracks/track.py b/unshackle/core/tracks/track.py index a315f655..daa1d9b4 100644 --- a/unshackle/core/tracks/track.py +++ b/unshackle/core/tracks/track.py @@ -439,6 +439,57 @@ class Track: self.path = target return target + def to_dict(self) -> dict[str, Any]: + """Serialise the track for export/import (identity/URL/descriptor/language). + + DRM is not serialised here; the export writer attaches the licensed DRM + keys. + Subclasses add their own codec/quality fields. + """ + data: dict[str, Any] = { + "type": self.__class__.__name__, + "id": self.id, + "url": self.url, + "language": str(self.language), + "is_original_lang": self.is_original_lang, + "descriptor": self.descriptor.name, + "needs_repack": self.needs_repack, + "name": self.name, + "edition": self.edition, + } + return data + + @staticmethod + def base_kwargs_from_dict(data: dict[str, Any]) -> dict[str, Any]: + """Build the shared Track constructor kwargs from a ``to_dict()`` payload. + + DRM is not reconstructed here — ``to_dict`` does not serialise it, and the import + flow attaches the licensed DRM + content keys separately. + """ + return { + "url": data["url"], + "language": data.get("language") or "und", + "is_original_lang": data.get("is_original_lang", False), + "descriptor": Track.Descriptor[data.get("descriptor", "URL")], + "needs_repack": data.get("needs_repack", False), + "name": data.get("name"), + "edition": data.get("edition") or None, + "id_": data.get("id"), + } + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> "Track": + """Reconstruct the correct Track subclass from a ``to_dict()`` payload.""" + from unshackle.core.tracks.audio import Audio + from unshackle.core.tracks.subtitle import Subtitle + from unshackle.core.tracks.video import Video + + track_type = data.get("type") + builders = {"Video": Video, "Audio": Audio, "Subtitle": Subtitle} + builder = builders.get(track_type) + if builder is None: + raise ValueError(f"Cannot reconstruct unsupported track type: {track_type!r}") + return builder.from_dict(data) + def get_track_name(self) -> Optional[str]: """Get the Track Name.""" return self.name diff --git a/unshackle/core/tracks/video.py b/unshackle/core/tracks/video.py index 39092e2b..e8678cc9 100644 --- a/unshackle/core/tracks/video.py +++ b/unshackle/core/tracks/video.py @@ -297,6 +297,37 @@ class Video(Track): self.needs_duration_fix = False self.dv_compatible_bitstream = dv_compatible_bitstream + def to_dict(self) -> dict[str, Any]: + data = super().to_dict() + data.update( + { + "codec": self.codec.name if self.codec else None, + "range": self.range.name if self.range else None, + "bitrate": self.bitrate, + "width": self.width, + "height": self.height, + "fps": str(self.fps) if self.fps else None, + "scan_type": self.scan_type.name if self.scan_type else None, + "dv_compatible_bitstream": self.dv_compatible_bitstream, + } + ) + return data + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> Video: + kwargs = Track.base_kwargs_from_dict(data) + return cls( + **kwargs, + codec=Video.Codec[data["codec"]] if data.get("codec") else None, + range_=Video.Range[data["range"]] if data.get("range") else None, + bitrate=data.get("bitrate"), + width=data.get("width"), + height=data.get("height"), + fps=data.get("fps"), + scan_type=Video.ScanType[data["scan_type"]] if data.get("scan_type") else None, + dv_compatible_bitstream=data.get("dv_compatible_bitstream", False), + ) + def __str__(self) -> str: return " | ".join( filter(