forked from DRMTalks/devine
refactor: Move dl command's download_track() to Track.download()
This commit is contained in:
parent
f510095bcf
commit
cae47017dc
|
@ -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."""
|
||||
|
|
|
@ -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()
|
||||
|
|
Loading…
Reference in New Issue