From b92708ef4567328cbb0000e5846196a40bbab76b Mon Sep 17 00:00:00 2001 From: rlaphoenix Date: Thu, 11 May 2023 22:17:41 +0100 Subject: [PATCH] Alter behaviour of --skip-dl to allow DRM licensing Most people used --skip-dl just to license the DRM pre-v1.3.0. Which makes sense, --skip-dl is otherwise a pointless feature. I've fixed it so that --skip-dl worked like before, allowing license calls, while still supporting the new per-segment features post-v1.3.0. Fixes #37 --- devine/commands/dl.py | 215 ++++++++++++++++++---------------- devine/core/manifests/dash.py | 5 + devine/core/manifests/hls.py | 5 + 3 files changed, 125 insertions(+), 100 deletions(-) diff --git a/devine/commands/dl.py b/devine/commands/dl.py index 51eb7c0..26be204 100644 --- a/devine/commands/dl.py +++ b/devine/commands/dl.py @@ -134,6 +134,7 @@ class dl: return dl(ctx, **kwargs) DL_POOL_STOP = Event() + DL_POOL_SKIP = Event() DRM_TABLE_LOCK = Lock() def __init__( @@ -458,75 +459,78 @@ class dl: download_table = Table.grid() download_table.add_row(selected_tracks) + dl_start_time = time.time() + if skip_dl: - self.log.info("Skipping Download...") - else: - dl_start_time = time.time() + self.DL_POOL_SKIP.set() - try: - with Live( - Padding( - download_table, - (1, 5) - ), - console=console, - refresh_per_second=5 - ): - with ThreadPoolExecutor(workers) as pool: - for download in futures.as_completed(( - pool.submit( - self.download_track, - service=service, - track=track, - prepare_drm=partial( - partial( - self.prepare_drm, - table=download_table - ), - track=track, - title=title, - certificate=partial( - service.get_widevine_service_certificate, - title=title, - track=track - ), - licence=partial( - service.get_widevine_license, - title=title, - track=track - ), - cdm_only=cdm_only, - vaults_only=vaults_only, - export=export - ), - progress=tracks_progress_callables[i] - ) - for i, track in enumerate(title.tracks) - )): - download.result() - except KeyboardInterrupt: - console.print(Padding( - ":x: Download Cancelled...", - (0, 5, 1, 5) - )) - return - except Exception as e: # noqa - error_messages = [ - ":x: Download Failed...", - " One of the download workers had an error!", - " See the error trace above for more information." - ] - if isinstance(e, subprocess.CalledProcessError): - # ignore process exceptions as proper error logs are already shown - error_messages.append(f" Process exit code: {e.returncode}") - else: - console.print_exception() - console.print(Padding( - Group(*error_messages), + try: + with Live( + Padding( + download_table, (1, 5) - )) - return + ), + console=console, + refresh_per_second=5 + ): + with ThreadPoolExecutor(workers) as pool: + for download in futures.as_completed(( + pool.submit( + self.download_track, + service=service, + track=track, + prepare_drm=partial( + partial( + self.prepare_drm, + table=download_table + ), + track=track, + title=title, + certificate=partial( + service.get_widevine_service_certificate, + title=title, + track=track + ), + licence=partial( + service.get_widevine_license, + title=title, + track=track + ), + cdm_only=cdm_only, + vaults_only=vaults_only, + export=export + ), + progress=tracks_progress_callables[i] + ) + for i, track in enumerate(title.tracks) + )): + download.result() + except KeyboardInterrupt: + console.print(Padding( + ":x: Download Cancelled...", + (0, 5, 1, 5) + )) + return + except Exception as e: # noqa + error_messages = [ + ":x: Download Failed...", + " One of the download workers had an error!", + " See the error trace above for more information." + ] + if isinstance(e, subprocess.CalledProcessError): + # ignore process exceptions as proper error logs are already shown + error_messages.append(f" Process exit code: {e.returncode}") + else: + console.print_exception() + console.print(Padding( + Group(*error_messages), + (1, 5) + )) + return + if skip_dl: + console.log("Skipped downloads as --skip-dl was used...") + else: dl_time = time_elapsed_since(dl_start_time) console.print(Padding( f"Track downloads finished in [progress.elapsed]{dl_time}[/]", @@ -806,6 +810,9 @@ class dl: prepare_drm: Callable, progress: partial ): + if self.DL_POOL_SKIP.is_set(): + progress(downloaded="[yellow]SKIPPING") + if self.DL_POOL_STOP.is_set(): progress(downloaded="[yellow]SKIPPED") return @@ -815,12 +822,6 @@ class dl: else: proxy = None - if config.directories.temp.is_file(): - self.log.error(f"Temp Directory '{config.directories.temp}' must be a Directory, not a file") - sys.exit(1) - - config.directories.temp.mkdir(parents=True, exist_ok=True) - save_path = config.directories.temp / f"{track.__class__.__name__}_{track.id}.mp4" if isinstance(track, Subtitle): save_path = save_path.with_suffix(f".{track.codec.extension}") @@ -841,11 +842,18 @@ class dl: if save_dir.exists() and save_dir.name.endswith("_segments"): shutil.rmtree(save_dir) - # Delete any pre-existing temp files matching this track. - # We can't re-use or continue downloading these tracks as they do not use a - # lock file. Or at least the majority don't. Even if they did I've encountered - # corruptions caused by sudden interruptions to the lock file. - cleanup() + if not self.DL_POOL_SKIP.is_set(): + if config.directories.temp.is_file(): + self.log.error(f"Temp Directory '{config.directories.temp}' must be a Directory, not a file") + sys.exit(1) + + config.directories.temp.mkdir(parents=True, exist_ok=True) + + # Delete any pre-existing temp files matching this track. + # We can't re-use or continue downloading these tracks as they do not use a + # lock file. Or at least the majority don't. Even if they did I've encountered + # corruptions caused by sudden interruptions to the lock file. + cleanup() try: if track.descriptor == track.Descriptor.M3U: @@ -854,6 +862,7 @@ class dl: save_path=save_path, save_dir=save_dir, stop_event=self.DL_POOL_STOP, + skip_event=self.DL_POOL_SKIP, progress=progress, session=service.session, proxy=proxy, @@ -865,6 +874,7 @@ class dl: save_path=save_path, save_dir=save_dir, stop_event=self.DL_POOL_STOP, + skip_event=self.DL_POOL_SKIP, progress=progress, session=service.session, proxy=proxy, @@ -893,21 +903,24 @@ class dl: else: drm = None - asyncio.run(aria2c( - uri=track.url, - out=save_path, - headers=service.session.headers, - proxy=proxy if track.needs_proxy else None, - progress=progress - )) + if self.DL_POOL_SKIP.is_set(): + progress(downloaded="[yellow]SKIPPED") + else: + asyncio.run(aria2c( + uri=track.url, + out=save_path, + headers=service.session.headers, + proxy=proxy if track.needs_proxy else None, + progress=progress + )) - track.path = save_path + track.path = save_path - if drm: - drm.decrypt(save_path) - track.drm = None - if callable(track.OnDecrypted): - track.OnDecrypted(track) + if drm: + drm.decrypt(save_path) + track.drm = None + if callable(track.OnDecrypted): + track.OnDecrypted(track) except KeyboardInterrupt: self.DL_POOL_STOP.set() progress(downloaded="[yellow]STOPPED") @@ -917,25 +930,27 @@ class dl: progress(downloaded="[red]FAILED") raise except (Exception, KeyboardInterrupt): - cleanup() + if not self.DL_POOL_SKIP.is_set(): + cleanup() raise if self.DL_POOL_STOP.is_set(): # we stopped during the download, let's exit return - if track.path.stat().st_size <= 3: # Empty UTF-8 BOM == 3 bytes - raise IOError( - "Download failed, the downloaded file is empty. " - f"This {'was' if track.needs_proxy else 'was not'} downloaded with a proxy." + - ( - " Perhaps you need to set `needs_proxy` as True to use the proxy for this track." - if not track.needs_proxy else "" + if not self.DL_POOL_SKIP.is_set(): + if track.path.stat().st_size <= 3: # Empty UTF-8 BOM == 3 bytes + raise IOError( + "Download failed, the downloaded file is empty. " + f"This {'was' if track.needs_proxy else 'was not'} downloaded with a proxy." + + ( + " Perhaps you need to set `needs_proxy` as True to use the proxy for this track." + if not track.needs_proxy else "" + ) ) - ) - if callable(track.OnDownloaded): - track.OnDownloaded(track) + if callable(track.OnDownloaded): + track.OnDownloaded(track) @staticmethod def get_profile(service: str) -> Optional[str]: diff --git a/devine/core/manifests/dash.py b/devine/core/manifests/dash.py index 8f4a5b4..2ced4c0 100644 --- a/devine/core/manifests/dash.py +++ b/devine/core/manifests/dash.py @@ -283,6 +283,7 @@ class DASH: save_path: Path, save_dir: Path, stop_event: Event, + skip_event: Event, progress: partial, session: Optional[Session] = None, proxy: Optional[str] = None, @@ -458,6 +459,10 @@ class DASH: else: drm = None + if skip_event.is_set(): + progress(downloaded="[yellow]SKIPPED") + return + def download_segment(filename: str, segment: tuple[str, Optional[str]]) -> int: if stop_event.is_set(): # the track already started downloading, but another failed or was stopped diff --git a/devine/core/manifests/hls.py b/devine/core/manifests/hls.py index cb53175..e9d7ad3 100644 --- a/devine/core/manifests/hls.py +++ b/devine/core/manifests/hls.py @@ -184,6 +184,7 @@ class HLS: save_path: Path, save_dir: Path, stop_event: Event, + skip_event: Event, progress: partial, session: Optional[Session] = None, proxy: Optional[str] = None, @@ -280,6 +281,10 @@ class HLS: finally: segment_key.put(newest_segment_key) + if skip_event.is_set(): + progress(downloaded="[yellow]SKIPPING") + return 0 + if not segment.uri.startswith(segment.base_uri): segment.uri = segment.base_uri + segment.uri