From cae47017dc3468b71c844036f77e680c9f351027 Mon Sep 17 00:00:00 2001 From: rlaphoenix Date: Sat, 2 Mar 2024 15:08:22 +0000 Subject: [PATCH] refactor: Move dl command's download_track() to Track.download() --- devine/commands/dl.py | 169 ++---------------------------------- devine/core/tracks/track.py | 168 ++++++++++++++++++++++++++++++++++- 2 files changed, 170 insertions(+), 167 deletions(-) diff --git a/devine/commands/dl.py b/devine/commands/dl.py index 3b36c71..9c55d07 100644 --- a/devine/commands/dl.py +++ b/devine/commands/dl.py @@ -40,18 +40,16 @@ from rich.tree import Tree from devine.core.config import config from devine.core.console import console -from devine.core.constants import DOWNLOAD_CANCELLED, DOWNLOAD_LICENCE_ONLY, AnyTrack, context_settings +from devine.core.constants import DOWNLOAD_LICENCE_ONLY, AnyTrack, context_settings from devine.core.credential import Credential -from devine.core.downloaders import downloader from devine.core.drm import DRM_T, Widevine -from devine.core.manifests import DASH, HLS from devine.core.proxies import Basic, Hola, NordVPN from devine.core.service import Service from devine.core.services import Services from devine.core.titles import Movie, Song, Title_T from devine.core.titles.episode import Episode from devine.core.tracks import Audio, Subtitle, Tracks, Video -from devine.core.utilities import get_binary_path, is_close_match, time_elapsed_since, try_ensure_utf8 +from devine.core.utilities import get_binary_path, is_close_match, time_elapsed_since from devine.core.utils.click_types import LANGUAGE_RANGE, QUALITY_LIST, SEASON_RANGE, ContextData from devine.core.utils.collections import merge_dict from devine.core.utils.subprocess import ffprobe @@ -376,13 +374,7 @@ class dl: sys.exit(1) video_languages = v_lang or lang - if ( - (v_lang and "all" not in v_lang) or - (lang and "all" not in lang and any( - x.language != title.tracks.videos[0].language - for x in title.tracks.videos - )) - ): + if video_languages and "all" not in video_languages: title.tracks.videos = title.tracks.by_language(title.tracks.videos, video_languages) if not title.tracks.videos: self.log.error(f"There's no {video_languages} Video Track...") @@ -476,9 +468,8 @@ class dl: with ThreadPoolExecutor(workers) as pool: for download in futures.as_completed(( pool.submit( - self.download_track, - service=service, - track=track, + track.download, + session=service.session, prepare_drm=partial( partial( self.prepare_drm, @@ -795,156 +786,6 @@ class dl: keys[str(title)][str(track)].update(drm.content_keys) export.write_text(jsonpickle.dumps(keys, indent=4), encoding="utf8") - def download_track( - self, - service: Service, - track: AnyTrack, - prepare_drm: Callable, - progress: partial - ): - if DOWNLOAD_LICENCE_ONLY.is_set(): - progress(downloaded="[yellow]SKIPPING") - - if DOWNLOAD_CANCELLED.is_set(): - progress(downloaded="[yellow]CANCELLED") - return - - proxy = next(iter(service.session.proxies.values()), None) - - 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}") - - if track.descriptor != track.Descriptor.URL: - save_dir = save_path.with_name(save_path.name + "_segments") - else: - save_dir = save_path.parent - - def cleanup(): - # track file (e.g., "foo.mp4") - save_path.unlink(missing_ok=True) - # aria2c control file (e.g., "foo.mp4.aria2") - save_path.with_suffix(f"{save_path.suffix}.aria2").unlink(missing_ok=True) - if save_dir.exists() and save_dir.name.endswith("_segments"): - shutil.rmtree(save_dir) - - if not DOWNLOAD_LICENCE_ONLY.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.HLS: - HLS.download_track( - track=track, - save_path=save_path, - save_dir=save_dir, - progress=progress, - session=service.session, - proxy=proxy, - license_widevine=prepare_drm - ) - elif track.descriptor == track.Descriptor.DASH: - DASH.download_track( - track=track, - save_path=save_path, - save_dir=save_dir, - progress=progress, - session=service.session, - proxy=proxy, - license_widevine=prepare_drm - ) - elif track.descriptor == track.Descriptor.URL: - try: - if not track.drm and isinstance(track, (Video, Audio)): - # the service might not have explicitly defined the `drm` property - # try find widevine DRM information from the init data of URL - try: - track.drm = [Widevine.from_track(track, service.session)] - except Widevine.Exceptions.PSSHNotFound: - # it might not have Widevine DRM, or might not have found the PSSH - self.log.warning("No Widevine PSSH was found for this track, is it DRM free?") - - if track.drm: - track_kid = track.get_key_id(session=service.session) - drm = track.drm[0] # just use the first supported DRM system for now - if isinstance(drm, Widevine): - # license and grab content keys - if not prepare_drm: - raise ValueError("prepare_drm func must be supplied to use Widevine DRM") - progress(downloaded="LICENSING") - prepare_drm(drm, track_kid=track_kid) - progress(downloaded="[yellow]LICENSED") - else: - drm = None - - if DOWNLOAD_LICENCE_ONLY.is_set(): - progress(downloaded="[yellow]SKIPPED") - else: - for status_update in downloader( - urls=track.url, - output_dir=save_path.parent, - filename=save_path.name, - headers=service.session.headers, - cookies=service.session.cookies, - proxy=proxy - ): - file_downloaded = status_update.get("file_downloaded") - if not file_downloaded: - progress(**status_update) - - track.path = save_path - if callable(track.OnDownloaded): - track.OnDownloaded() - - if drm: - progress(downloaded="Decrypting", completed=0, total=100) - drm.decrypt(save_path) - track.drm = None - if callable(track.OnDecrypted): - track.OnDecrypted(drm) - progress(downloaded="Decrypted", completed=100) - - if isinstance(track, Subtitle) and \ - track.codec not in (Subtitle.Codec.fVTT, Subtitle.Codec.fTTML): - track_data = track.path.read_bytes() - track_data = try_ensure_utf8(track_data) - track_data = track_data.decode("utf8"). \ - replace("‎", html.unescape("‎")). \ - replace("‏", html.unescape("‏")). \ - encode("utf8") - track.path.write_bytes(track_data) - - progress(downloaded="Downloaded") - except KeyboardInterrupt: - DOWNLOAD_CANCELLED.set() - progress(downloaded="[yellow]CANCELLED") - raise - except Exception: - DOWNLOAD_CANCELLED.set() - progress(downloaded="[red]FAILED") - raise - except (Exception, KeyboardInterrupt): - if not DOWNLOAD_LICENCE_ONLY.is_set(): - cleanup() - raise - - if DOWNLOAD_CANCELLED.is_set(): - # we stopped during the download, let's exit - return - - if not DOWNLOAD_LICENCE_ONLY.is_set(): - if track.path.stat().st_size <= 3: # Empty UTF-8 BOM == 3 bytes - raise IOError("Download failed, the downloaded file is empty.") - @staticmethod def get_profile(service: str) -> Optional[str]: """Get profile for Service from config.""" diff --git a/devine/core/tracks/track.py b/devine/core/tracks/track.py index 39e3c27..0c75c22 100644 --- a/devine/core/tracks/track.py +++ b/devine/core/tracks/track.py @@ -1,9 +1,12 @@ import base64 +import html +import logging import re import shutil import subprocess from copy import copy from enum import Enum +from functools import partial from pathlib import Path from typing import Any, Callable, Iterable, Optional, Union from uuid import UUID @@ -13,9 +16,11 @@ import m3u8 import requests from langcodes import Language -from devine.core.constants import TERRITORY_MAP -from devine.core.drm import DRM_T -from devine.core.utilities import get_binary_path, get_boxes +from devine.core.config import config +from devine.core.constants import DOWNLOAD_CANCELLED, DOWNLOAD_LICENCE_ONLY, TERRITORY_MAP +from devine.core.downloaders import downloader +from devine.core.drm import DRM_T, Widevine +from devine.core.utilities import get_binary_path, get_boxes, try_ensure_utf8 from devine.core.utils.subprocess import ffprobe @@ -109,6 +114,163 @@ class Track: def __eq__(self, other: Any) -> bool: return isinstance(other, Track) and self.id == other.id + def download( + self, + session: requests.Session, + prepare_drm: partial, + progress: Optional[partial] = None + ): + """Download and optionally Decrypt this Track.""" + from devine.core.manifests import DASH, HLS + + if DOWNLOAD_LICENCE_ONLY.is_set(): + progress(downloaded="[yellow]SKIPPING") + + if DOWNLOAD_CANCELLED.is_set(): + progress(downloaded="[yellow]SKIPPED") + return + + log = logging.getLogger("track") + + proxy = next(iter(session.proxies.values()), None) + + track_type = self.__class__.__name__ + save_path = config.directories.temp / f"{track_type}_{self.id}.mp4" + if track_type == "Subtitle": + save_path = save_path.with_suffix(f".{self.codec.extension}") + + if self.descriptor != self.Descriptor.URL: + save_dir = save_path.with_name(save_path.name + "_segments") + else: + save_dir = save_path.parent + + def cleanup(): + # track file (e.g., "foo.mp4") + save_path.unlink(missing_ok=True) + # aria2c control file (e.g., "foo.mp4.aria2" or "foo.mp4.aria2__temp") + save_path.with_suffix(f"{save_path.suffix}.aria2").unlink(missing_ok=True) + save_path.with_suffix(f"{save_path.suffix}.aria2__temp").unlink(missing_ok=True) + if save_dir.exists() and save_dir.name.endswith("_segments"): + shutil.rmtree(save_dir) + + if not DOWNLOAD_LICENCE_ONLY.is_set(): + if config.directories.temp.is_file(): + raise ValueError(f"Temp Directory '{config.directories.temp}' must be a Directory, not a file") + + 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 self.descriptor == self.Descriptor.HLS: + HLS.download_track( + track=self, + save_path=save_path, + save_dir=save_dir, + progress=progress, + session=session, + proxy=proxy, + license_widevine=prepare_drm + ) + elif self.descriptor == self.Descriptor.DASH: + DASH.download_track( + track=self, + save_path=save_path, + save_dir=save_dir, + progress=progress, + session=session, + proxy=proxy, + license_widevine=prepare_drm + ) + elif self.descriptor == self.Descriptor.URL: + try: + if not self.drm and track_type in ("Video", "Audio"): + # the service might not have explicitly defined the `drm` property + # try find widevine DRM information from the init data of URL + try: + self.drm = [Widevine.from_track(self, session)] + except Widevine.Exceptions.PSSHNotFound: + # it might not have Widevine DRM, or might not have found the PSSH + log.warning("No Widevine PSSH was found for this track, is it DRM free?") + + if self.drm: + track_kid = self.get_key_id(session=session) + drm = self.drm[0] # just use the first supported DRM system for now + if isinstance(drm, Widevine): + # license and grab content keys + if not prepare_drm: + raise ValueError("prepare_drm func must be supplied to use Widevine DRM") + progress(downloaded="LICENSING") + prepare_drm(drm, track_kid=track_kid) + progress(downloaded="[yellow]LICENSED") + else: + drm = None + + if DOWNLOAD_LICENCE_ONLY.is_set(): + progress(downloaded="[yellow]SKIPPED") + else: + for status_update in downloader( + urls=self.url, + output_dir=save_path.parent, + filename=save_path.name, + headers=session.headers, + cookies=session.cookies, + proxy=proxy + ): + file_downloaded = status_update.get("file_downloaded") + if not file_downloaded: + progress(**status_update) + + self.path = save_path + if callable(self.OnDownloaded): + self.OnDownloaded() + + if drm: + progress(downloaded="Decrypting", completed=0, total=100) + drm.decrypt(save_path) + self.drm = None + if callable(self.OnDecrypted): + self.OnDecrypted(drm) + progress(downloaded="Decrypted", completed=100) + + if track_type == "Subtitle" and self.codec.name not in ("fVTT", "fTTML"): + track_data = self.path.read_bytes() + track_data = try_ensure_utf8(track_data) + track_data = track_data.decode("utf8"). \ + replace("‎", html.unescape("‎")). \ + replace("‏", html.unescape("‏")). \ + encode("utf8") + self.path.write_bytes(track_data) + + progress(downloaded="Downloaded") + except KeyboardInterrupt: + DOWNLOAD_CANCELLED.set() + progress(downloaded="[yellow]CANCELLED") + raise + except Exception: + DOWNLOAD_CANCELLED.set() + progress(downloaded="[red]FAILED") + raise + except (Exception, KeyboardInterrupt): + if not DOWNLOAD_LICENCE_ONLY.is_set(): + cleanup() + raise + + if DOWNLOAD_CANCELLED.is_set(): + # we stopped during the download, let's exit + return + + if not DOWNLOAD_LICENCE_ONLY.is_set(): + if self.path.stat().st_size <= 3: # Empty UTF-8 BOM == 3 bytes + raise IOError("Download failed, the downloaded file is empty.") + + if callable(self.OnDownloaded): + self.OnDownloaded(self) + def delete(self) -> None: if self.path: self.path.unlink()