From cae47017dc3468b71c844036f77e680c9f351027 Mon Sep 17 00:00:00 2001
From: rlaphoenix <rlaphoenix@pm.me>
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()