diff --git a/unshackle/commands/dl.py b/unshackle/commands/dl.py index 46a6dd9..86f56ad 100644 --- a/unshackle/commands/dl.py +++ b/unshackle/commands/dl.py @@ -482,7 +482,12 @@ class dl: @click.option( "--skip-dl", is_flag=True, default=False, help="Skip downloading while still retrieving the decryption keys." ) - @click.option("--export", type=Path, help="Export Decryption Keys as you obtain them to a JSON file.") + @click.option( + "--export", + is_flag=True, + default=False, + help="Export track info and decryption keys to a JSON file in the exports directory.", + ) @click.option( "--cdm-only/--vaults-only", is_flag=True, @@ -546,6 +551,7 @@ class dl: return dl(ctx, **kwargs) DRM_TABLE_LOCK = Lock() + EXPORT_LOCK = Lock() def __init__( self, @@ -1022,7 +1028,7 @@ class dl: list_: bool, list_titles: bool, skip_dl: bool, - export: Optional[Path], + export: bool, cdm_only: Optional[bool], no_proxy: bool, no_folder: bool, @@ -1045,6 +1051,12 @@ class dl: if skip_dl: DOWNLOAD_LICENCE_ONLY.set() + if export: + config.directories.exports.mkdir(parents=True, exist_ok=True) + export_path = config.directories.exports / f"export_{self.service}_{int(time.time())}.json" + else: + export_path = None + # Parse bitrate range options vbitrate_min, vbitrate_max = None, None if vbitrate_range: @@ -1998,7 +2010,7 @@ class dl: kept_tracks.extend(title.tracks.chapters) kept_tracks.extend(title.tracks.attachments) - title.tracks = Tracks(kept_tracks) + title.tracks = Tracks(kept_tracks, manifest_url=title.tracks.manifest_url) selected_tracks, tracks_progress_callables = title.tracks.tree(add_progress=True) @@ -2063,7 +2075,7 @@ class dl: ), cdm_only=cdm_only, vaults_only=vaults_only, - export=export, + export=export_path, ), cdm=self.cdm, max_workers=workers, @@ -2593,6 +2605,55 @@ 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.""" + with self.EXPORT_LOCK: + keys = {} + if export.is_file(): + keys = jsonpickle.loads(export.read_text(encoding="utf8")) or {} + if str(title) not in keys: + keys[str(title)] = {} + + title_data = keys[str(title)] + + if title.tracks.manifest_url and "manifest" not in title_data: + title_data["manifest"] = title.tracks.manifest_url + + if isinstance(track, Video): + section = "video" + elif isinstance(track, Audio): + section = "audio" + else: + section = "other" + + if section not in title_data: + title_data[section] = {} + if str(track) not in title_data[section]: + title_data[section][str(track)] = {} + + track_data = title_data[section][str(track)] + track_data["url"] = track.url + track_data["descriptor"] = track.descriptor.name + + if "keys" not in track_data: + track_data["keys"] = {} + for kid, key in drm.content_keys.items(): + track_data["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 + + 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") + def prepare_drm( self, drm: DRM_T, @@ -2851,24 +2912,7 @@ class dl: table.add_row(cek_tree) if export: - keys = {} - if export.is_file(): - keys = jsonpickle.loads(export.read_text(encoding="utf8")) or {} - if str(title) not in keys: - keys[str(title)] = {} - if str(track) not in keys[str(title)]: - keys[str(title)][str(track)] = {} - - track_data = keys[str(title)][str(track)] - track_data["url"] = track.url - track_data["descriptor"] = track.descriptor.name - - if "keys" not in track_data: - track_data["keys"] = {} - for kid, key in drm.content_keys.items(): - track_data["keys"][kid.hex] = key - - export.write_text(jsonpickle.dumps(keys, indent=4), encoding="utf8") + self._write_export(export, title, track, drm) elif isinstance(drm, PlayReady): if self.debug_logger: @@ -3016,24 +3060,7 @@ class dl: table.add_row(cek_tree) if export: - keys = {} - if export.is_file(): - keys = jsonpickle.loads(export.read_text(encoding="utf8")) or {} - if str(title) not in keys: - keys[str(title)] = {} - if str(track) not in keys[str(title)]: - keys[str(title)][str(track)] = {} - - track_data = keys[str(title)][str(track)] - track_data["url"] = track.url - track_data["descriptor"] = track.descriptor.name - - if "keys" not in track_data: - track_data["keys"] = {} - for kid, key in drm.content_keys.items(): - track_data["keys"][kid.hex] = key - - export.write_text(jsonpickle.dumps(keys, indent=4), encoding="utf8") + self._write_export(export, title, track, drm) elif isinstance(drm, MonaLisa): with self.DRM_TABLE_LOCK: diff --git a/unshackle/core/config.py b/unshackle/core/config.py index 64a1585..06bf787 100644 --- a/unshackle/core/config.py +++ b/unshackle/core/config.py @@ -26,6 +26,7 @@ class Config: cache = data / "cache" cookies = data / "cookies" logs = data / "logs" + exports = data / "exports" wvds = data / "WVDs" prds = data / "PRDs" dcsl = data / "DCSL" diff --git a/unshackle/core/manifests/dash.py b/unshackle/core/manifests/dash.py index a09d705..ac63164 100644 --- a/unshackle/core/manifests/dash.py +++ b/unshackle/core/manifests/dash.py @@ -247,6 +247,7 @@ class DASH: # only get tracks from the first main-content period break + tracks.manifest_url = self.url return tracks @staticmethod diff --git a/unshackle/core/manifests/hls.py b/unshackle/core/manifests/hls.py index ddab100..08ad646 100644 --- a/unshackle/core/manifests/hls.py +++ b/unshackle/core/manifests/hls.py @@ -36,7 +36,9 @@ from unshackle.core.utilities import get_debug_logger, get_extension, is_close_m class HLS: - def __init__(self, manifest: M3U8, session: Optional[Union[Session, RnetSession]] = None): + def __init__( + self, manifest: M3U8, session: Optional[Union[Session, RnetSession]] = None, url: Optional[str] = None + ): if not manifest: raise ValueError("HLS manifest must be provided.") if not isinstance(manifest, M3U8): @@ -46,6 +48,7 @@ class HLS: self.manifest = manifest self.session = session or Session() + self.url = url @classmethod def from_url(cls, url: str, session: Optional[Union[Session, RnetSession]] = None, **args: Any) -> HLS: @@ -75,7 +78,7 @@ class HLS: master = m3u8.loads(content, uri=url) - return cls(master, session) + return cls(master, session, url=url) @classmethod def from_text(cls, text: str, url: str) -> HLS: @@ -255,6 +258,8 @@ class HLS: except Exception: pass + if self.url: + tracks.manifest_url = self.url return tracks @staticmethod diff --git a/unshackle/core/manifests/ism.py b/unshackle/core/manifests/ism.py index 875fd5e..778a625 100644 --- a/unshackle/core/manifests/ism.py +++ b/unshackle/core/manifests/ism.py @@ -219,6 +219,7 @@ class ISM: data=data, ) ) + tracks.manifest_url = self.url return tracks @staticmethod diff --git a/unshackle/core/manifests/m3u8.py b/unshackle/core/manifests/m3u8.py index 2f39f84..336f5d7 100644 --- a/unshackle/core/manifests/m3u8.py +++ b/unshackle/core/manifests/m3u8.py @@ -17,9 +17,10 @@ def parse( language: str, *, session: Optional[Union[Session, RnetSession]] = None, + url: Optional[str] = None, ) -> Tracks: """Parse a variant playlist to ``Tracks`` with basic information, defer DRM loading.""" - tracks = HLS(master, session=session).to_tracks(language) + tracks = HLS(master, session=session, url=url).to_tracks(language) bool(master.session_keys or HLS.parse_session_data_keys(master, session or Session())) diff --git a/unshackle/core/service.py b/unshackle/core/service.py index fa3c346..f15c2d4 100644 --- a/unshackle/core/service.py +++ b/unshackle/core/service.py @@ -276,6 +276,8 @@ class Service(metaclass=ABCMeta): raise if first: all_tracks.add(hdr_tracks, warn_only=True) + if hdr_tracks.manifest_url and not all_tracks.manifest_url: + all_tracks.manifest_url = hdr_tracks.manifest_url first = False else: for video in hdr_tracks.videos: @@ -299,6 +301,8 @@ class Service(metaclass=ABCMeta): raise if first: all_tracks.add(tracks, warn_only=True) + if tracks.manifest_url and not all_tracks.manifest_url: + all_tracks.manifest_url = tracks.manifest_url first = False else: for video in tracks.videos: diff --git a/unshackle/core/tracks/tracks.py b/unshackle/core/tracks/tracks.py index 7c6cc7f..412b2a6 100644 --- a/unshackle/core/tracks/tracks.py +++ b/unshackle/core/tracks/tracks.py @@ -39,12 +39,14 @@ class Tracks: *args: Union[ Tracks, Sequence[Union[AnyTrack, Chapter, Chapters, Attachment]], Track, Chapter, Chapters, Attachment ], + manifest_url: Optional[str] = None, ): self.videos: list[Video] = [] self.audio: list[Audio] = [] self.subtitles: list[Subtitle] = [] self.chapters = Chapters() self.attachments: list[Attachment] = [] + self.manifest_url: Optional[str] = manifest_url if args: self.add(args) @@ -195,6 +197,8 @@ class Tracks: ) -> None: """Add a provided track to its appropriate array and ensuring it's not a duplicate.""" if isinstance(tracks, Tracks): + if tracks.manifest_url and not self.manifest_url: + self.manifest_url = tracks.manifest_url tracks = [*list(tracks), *tracks.chapters, *tracks.attachments] duplicates = 0 @@ -302,6 +306,16 @@ class Tracks: def select_subtitles(self, x: Callable[[Subtitle], bool]) -> None: self.subtitles = list(filter(x, self.subtitles)) + def filter(self, predicate: Callable[[AnyTrack], bool]) -> Tracks: + """Return a new Tracks with tracks filtered by predicate, preserving metadata.""" + new_tracks = Tracks(manifest_url=self.manifest_url) + new_tracks.videos = [t for t in self.videos if predicate(t)] + new_tracks.audio = [t for t in self.audio if predicate(t)] + new_tracks.subtitles = [t for t in self.subtitles if predicate(t)] + new_tracks.chapters = self.chapters + new_tracks.attachments = list(self.attachments) + return new_tracks + def select_hybrid(self, tracks, quality): # Prefer HDR10+ over HDR10 as the base layer (preserves dynamic metadata) base_ranges = (Video.Range.HDR10P, Video.Range.HDR10)