ninjajiraiya
/
devine
forked from DRMTalks/devine
1
0
Fork 0
Code Pull Requests Activity

Move Download Cancel/Skip Events to constants

This commit is contained in:
rlaphoenix 2024-01-09 11:41:57 +00:00
parent ce457df151
commit fa3cee11b7
4 changed files with 41 additions and 61 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

32
devine/commands/dl.py
View File

@ -16,7 +16,7 @@ from functools import partial
from http.cookiejar import MozillaCookieJar from http.cookiejar import MozillaCookieJar
from itertools import zip_longest from itertools import zip_longest
from pathlib import Path from pathlib import Path
from threading import Event, Lock from threading import Lock
from typing import Any, Callable, Optional from typing import Any, Callable, Optional
from uuid import UUID from uuid import UUID
@ -41,7 +41,7 @@ from rich.tree import Tree
from devine.core.config import config from devine.core.config import config
from devine.core.console import console from devine.core.console import console
from devine.core.constants import AnyTrack, context_settings from devine.core.constants import DOWNLOAD_CANCELLED, DOWNLOAD_LICENCE_ONLY, AnyTrack, context_settings
from devine.core.credential import Credential from devine.core.credential import Credential
from devine.core.downloaders import downloader from devine.core.downloaders import downloader
from devine.core.drm import DRM_T, Widevine from devine.core.drm import DRM_T, Widevine
@ -138,8 +138,6 @@ class dl:
def cli(ctx: click.Context, **kwargs: Any) -> dl: def cli(ctx: click.Context, **kwargs: Any) -> dl:
return dl(ctx, **kwargs) return dl(ctx, **kwargs)
DL_POOL_STOP = Event()
DL_POOL_SKIP = Event()
DRM_TABLE_LOCK = Lock() DRM_TABLE_LOCK = Lock()
def __init__( def __init__(
@ -466,7 +464,7 @@ class dl:
dl_start_time = time.time() dl_start_time = time.time()
if skip_dl: if skip_dl:
self.DL_POOL_SKIP.set() DOWNLOAD_LICENCE_ONLY.set()
try: try:
with Live( with Live(
@ -827,11 +825,11 @@ class dl:
prepare_drm: Callable, prepare_drm: Callable,
progress: partial progress: partial
): ):
if self.DL_POOL_SKIP.is_set(): if DOWNLOAD_LICENCE_ONLY.is_set():
progress(downloaded="[yellow]SKIPPING") progress(downloaded="[yellow]SKIPPING")
if self.DL_POOL_STOP.is_set(): if DOWNLOAD_CANCELLED.is_set():
progress(downloaded="[yellow]SKIPPED") progress(downloaded="[yellow]CANCELLED")
return return
proxy = next(iter(service.session.proxies.values()), None) proxy = next(iter(service.session.proxies.values()), None)
@ -856,7 +854,7 @@ class dl:
if save_dir.exists() and save_dir.name.endswith("_segments"): if save_dir.exists() and save_dir.name.endswith("_segments"):
shutil.rmtree(save_dir) shutil.rmtree(save_dir)
if not self.DL_POOL_SKIP.is_set(): if not DOWNLOAD_LICENCE_ONLY.is_set():
if config.directories.temp.is_file(): if config.directories.temp.is_file():
self.log.error(f"Temp Directory '{config.directories.temp}' must be a Directory, not a file") self.log.error(f"Temp Directory '{config.directories.temp}' must be a Directory, not a file")
sys.exit(1) sys.exit(1)
@ -875,8 +873,6 @@ class dl:
track=track, track=track,
save_path=save_path, save_path=save_path,
save_dir=save_dir, save_dir=save_dir,
stop_event=self.DL_POOL_STOP,
skip_event=self.DL_POOL_SKIP,
progress=progress, progress=progress,
session=service.session, session=service.session,
proxy=proxy, proxy=proxy,
@ -887,8 +883,6 @@ class dl:
track=track, track=track,
save_path=save_path, save_path=save_path,
save_dir=save_dir, save_dir=save_dir,
stop_event=self.DL_POOL_STOP,
skip_event=self.DL_POOL_SKIP,
progress=progress, progress=progress,
session=service.session, session=service.session,
proxy=proxy, proxy=proxy,
@ -919,7 +913,7 @@ class dl:
else: else:
drm = None drm = None
if self.DL_POOL_SKIP.is_set(): if DOWNLOAD_LICENCE_ONLY.is_set():
progress(downloaded="[yellow]SKIPPED") progress(downloaded="[yellow]SKIPPED")
else: else:
downloader( downloader(
@ -948,23 +942,23 @@ class dl:
progress(downloaded="Downloaded") progress(downloaded="Downloaded")
except KeyboardInterrupt: except KeyboardInterrupt:
self.DL_POOL_STOP.set() DOWNLOAD_CANCELLED.set()
progress(downloaded="[yellow]CANCELLED") progress(downloaded="[yellow]CANCELLED")
raise raise
except Exception: except Exception:
self.DL_POOL_STOP.set() DOWNLOAD_CANCELLED.set()
progress(downloaded="[red]FAILED") progress(downloaded="[red]FAILED")
raise raise
except (Exception, KeyboardInterrupt): except (Exception, KeyboardInterrupt):
if not self.DL_POOL_SKIP.is_set(): if not DOWNLOAD_LICENCE_ONLY.is_set():
cleanup() cleanup()
raise raise
if self.DL_POOL_STOP.is_set(): if DOWNLOAD_CANCELLED.is_set():
# we stopped during the download, let's exit # we stopped during the download, let's exit
return return
if not self.DL_POOL_SKIP.is_set(): if not DOWNLOAD_LICENCE_ONLY.is_set():
if track.path.stat().st_size <= 3: # Empty UTF-8 BOM == 3 bytes if track.path.stat().st_size <= 3: # Empty UTF-8 BOM == 3 bytes
raise IOError("Download failed, the downloaded file is empty.") raise IOError("Download failed, the downloaded file is empty.")

4
devine/core/constants.py
View File

@ -1,5 +1,9 @@
from threading import Event
from typing import TypeVar, Union from typing import TypeVar, Union
DOWNLOAD_CANCELLED = Event()
DOWNLOAD_LICENCE_ONLY = Event()
DRM_SORT_MAP = ["ClearKey", "Widevine"] DRM_SORT_MAP = ["ClearKey", "Widevine"]
LANGUAGE_MUX_MAP = { LANGUAGE_MUX_MAP = {
# List of language tags that cannot be used by mkvmerge and need replacements. # List of language tags that cannot be used by mkvmerge and need replacements.

29
devine/core/manifests/dash.py
View File

@ -12,7 +12,6 @@ from copy import copy
from functools import partial from functools import partial
from hashlib import md5 from hashlib import md5
from pathlib import Path from pathlib import Path
from threading import Event
from typing import Any, Callable, MutableMapping, Optional, Union from typing import Any, Callable, MutableMapping, Optional, Union
from urllib.parse import urljoin, urlparse from urllib.parse import urljoin, urlparse
from uuid import UUID from uuid import UUID
@ -26,7 +25,7 @@ from requests import Session
from requests.cookies import RequestsCookieJar from requests.cookies import RequestsCookieJar
from rich import filesize from rich import filesize
from devine.core.constants import AnyTrack from devine.core.constants import DOWNLOAD_CANCELLED, DOWNLOAD_LICENCE_ONLY, AnyTrack
from devine.core.downloaders import downloader from devine.core.downloaders import downloader
from devine.core.downloaders import requests as requests_downloader from devine.core.downloaders import requests as requests_downloader
from devine.core.drm import Widevine from devine.core.drm import Widevine
@ -225,8 +224,6 @@ class DASH:
track: AnyTrack, track: AnyTrack,
save_path: Path, save_path: Path,
save_dir: Path, save_dir: Path,
stop_event: Event,
skip_event: Event,
progress: partial, progress: partial,
session: Optional[Session] = None, session: Optional[Session] = None,
proxy: Optional[str] = None, proxy: Optional[str] = None,
@ -401,13 +398,13 @@ class DASH:
license_widevine(drm, track_kid=track_kid) license_widevine(drm, track_kid=track_kid)
progress(downloaded="[yellow]LICENSED") progress(downloaded="[yellow]LICENSED")
except Exception: # noqa except Exception: # noqa
stop_event.set() # skip pending track downloads DOWNLOAD_CANCELLED.set() # skip pending track downloads
progress(downloaded="[red]FAILED") progress(downloaded="[red]FAILED")
raise raise
else: else:
drm = None drm = None
if skip_event.is_set(): if DOWNLOAD_LICENCE_ONLY.is_set():
progress(downloaded="[yellow]SKIPPED") progress(downloaded="[yellow]SKIPPED")
return return
@ -427,15 +424,14 @@ class DASH:
proxy=proxy, proxy=proxy,
headers=session.headers, headers=session.headers,
cookies=session.cookies, cookies=session.cookies,
bytes_range=bytes_range, bytes_range=bytes_range
stop_event=stop_event
) )
for n, (url, bytes_range) in enumerate(segments) for n, (url, bytes_range) in enumerate(segments)
))): ))):
try: try:
download_size = download.result() download_size = download.result()
except KeyboardInterrupt: except KeyboardInterrupt:
stop_event.set() # skip pending track downloads DOWNLOAD_CANCELLED.set() # skip pending track downloads
progress(downloaded="[yellow]CANCELLING") progress(downloaded="[yellow]CANCELLING")
pool.shutdown(wait=True, cancel_futures=True) pool.shutdown(wait=True, cancel_futures=True)
progress(downloaded="[yellow]CANCELLED") progress(downloaded="[yellow]CANCELLED")
@ -443,7 +439,7 @@ class DASH:
# the pool is already shut down, so exiting loop is fine # the pool is already shut down, so exiting loop is fine
raise raise
except Exception: except Exception:
stop_event.set() # skip pending track downloads DOWNLOAD_CANCELLED.set() # skip pending track downloads
progress(downloaded="[red]FAILING") progress(downloaded="[red]FAILING")
pool.shutdown(wait=True, cancel_futures=True) pool.shutdown(wait=True, cancel_futures=True)
progress(downloaded="[red]FAILED") progress(downloaded="[red]FAILED")
@ -501,8 +497,7 @@ class DASH:
proxy: Optional[str] = None, proxy: Optional[str] = None,
headers: Optional[MutableMapping[str, str | bytes]] = None, headers: Optional[MutableMapping[str, str | bytes]] = None,
cookies: Optional[Union[MutableMapping[str, str], RequestsCookieJar]] = None, cookies: Optional[Union[MutableMapping[str, str], RequestsCookieJar]] = None,
bytes_range: Optional[str] = None, bytes_range: Optional[str] = None
stop_event: Optional[Event] = None
) -> int: ) -> int:
""" """
Download a DASH Media Segment. Download a DASH Media Segment.
@ -518,12 +513,10 @@ class DASH:
will be resolved based on the URI among other parameters. Multiple cookies with will be resolved based on the URI among other parameters. Multiple cookies with
the same name but a different domain/path are resolved. the same name but a different domain/path are resolved.
bytes_range: Download only specific bytes of the Segment file using the Range header. bytes_range: Download only specific bytes of the Segment file using the Range header.
stop_event: Prematurely stop the Download from beginning. Useful if ran from
a Thread Pool. It will raise a KeyboardInterrupt if set.
Returns the file size of the downloaded Segment in bytes. Returns the file size of the downloaded Segment in bytes.
""" """
if stop_event and stop_event.is_set(): if DOWNLOAD_CANCELLED.is_set():
raise KeyboardInterrupt() raise KeyboardInterrupt()
attempts = 1 attempts = 1
@ -546,9 +539,9 @@ class DASH:
segmented=True segmented=True
) )
break break
except Exception as ee: except Exception as e:
if (stop_event and stop_event.is_set()) or attempts == 5: if DOWNLOAD_CANCELLED.is_set() or attempts == 5:
raise ee raise e
time.sleep(2) time.sleep(2)
attempts += 1 attempts += 1

37
devine/core/manifests/hls.py
View File

@ -10,7 +10,7 @@ from functools import partial
from hashlib import md5 from hashlib import md5
from pathlib import Path from pathlib import Path
from queue import Queue from queue import Queue
from threading import Event, Lock from threading import Lock
from typing import Any, Callable, Optional, Union from typing import Any, Callable, Optional, Union
from urllib.parse import urljoin from urllib.parse import urljoin
@ -23,7 +23,7 @@ from pywidevine.pssh import PSSH
from requests import Session from requests import Session
from rich import filesize from rich import filesize
from devine.core.constants import AnyTrack from devine.core.constants import DOWNLOAD_CANCELLED, DOWNLOAD_LICENCE_ONLY, AnyTrack
from devine.core.downloaders import downloader from devine.core.downloaders import downloader
from devine.core.downloaders import requests as requests_downloader from devine.core.downloaders import requests as requests_downloader
from devine.core.drm import DRM_T, ClearKey, Widevine from devine.core.drm import DRM_T, ClearKey, Widevine
@ -190,8 +190,6 @@ class HLS:
track: AnyTrack, track: AnyTrack,
save_path: Path, save_path: Path,
save_dir: Path, save_dir: Path,
stop_event: Event,
skip_event: Event,
progress: partial, progress: partial,
session: Optional[Session] = None, session: Optional[Session] = None,
proxy: Optional[str] = None, proxy: Optional[str] = None,
@ -231,7 +229,7 @@ class HLS:
license_widevine(session_drm) license_widevine(session_drm)
progress(downloaded="[yellow]LICENSED") progress(downloaded="[yellow]LICENSED")
except Exception: # noqa except Exception: # noqa
stop_event.set() # skip pending track downloads DOWNLOAD_CANCELLED.set() # skip pending track downloads
progress(downloaded="[red]FAILED") progress(downloaded="[red]FAILED")
raise raise
else: else:
@ -265,16 +263,14 @@ class HLS:
progress=progress, progress=progress,
license_widevine=license_widevine, license_widevine=license_widevine,
session=session, session=session,
proxy=proxy, proxy=proxy
stop_event=stop_event,
skip_event=skip_event
) )
for n, segment in enumerate(master.segments) for n, segment in enumerate(master.segments)
))): ))):
try: try:
download_size = download.result() download_size = download.result()
except KeyboardInterrupt: except KeyboardInterrupt:
stop_event.set() # skip pending track downloads DOWNLOAD_CANCELLED.set() # skip pending track downloads
progress(downloaded="[yellow]CANCELLING") progress(downloaded="[yellow]CANCELLING")
pool.shutdown(wait=True, cancel_futures=True) pool.shutdown(wait=True, cancel_futures=True)
progress(downloaded="[yellow]CANCELLED") progress(downloaded="[yellow]CANCELLED")
@ -282,7 +278,7 @@ class HLS:
# the pool is already shut down, so exiting loop is fine # the pool is already shut down, so exiting loop is fine
raise raise
except Exception as e: except Exception as e:
stop_event.set() # skip pending track downloads DOWNLOAD_CANCELLED.set() # skip pending track downloads
progress(downloaded="[red]FAILING") progress(downloaded="[red]FAILING")
pool.shutdown(wait=True, cancel_futures=True) pool.shutdown(wait=True, cancel_futures=True)
progress(downloaded="[red]FAILED") progress(downloaded="[red]FAILED")
@ -310,7 +306,7 @@ class HLS:
last_speed_refresh = now last_speed_refresh = now
download_sizes.clear() download_sizes.clear()
if skip_event.is_set(): if DOWNLOAD_LICENCE_ONLY.is_set():
return return
with open(save_path, "wb") as f: with open(save_path, "wb") as f:
@ -338,9 +334,7 @@ class HLS:
progress: partial, progress: partial,
license_widevine: Optional[Callable] = None, license_widevine: Optional[Callable] = None,
session: Optional[Session] = None, session: Optional[Session] = None,
proxy: Optional[str] = None, proxy: Optional[str] = None
stop_event: Optional[Event] = None,
skip_event: Optional[Event] = None
) -> int: ) -> int:
""" """
Download (and Decrypt) an HLS Media Segment. Download (and Decrypt) an HLS Media Segment.
@ -365,15 +359,10 @@ class HLS:
if the Segment's DRM uses Widevine. if the Segment's DRM uses Widevine.
proxy: Proxy URI to use when downloading the Segment file. proxy: Proxy URI to use when downloading the Segment file.
session: Python-Requests Session used when requesting init data. session: Python-Requests Session used when requesting init data.
stop_event: Prematurely stop the Download from beginning. Useful if ran from
a Thread Pool. It will raise a KeyboardInterrupt if set.
skip_event: Prematurely stop the Download from beginning. It returns with a
file size of -1 directly after DRM licensing occurs, even if it's DRM-free.
This is mainly for `--skip-dl` to allow licensing without downloading.
Returns the file size of the downloaded Segment in bytes. Returns the file size of the downloaded Segment in bytes.
""" """
if stop_event.is_set(): if DOWNLOAD_CANCELLED.is_set():
raise KeyboardInterrupt() raise KeyboardInterrupt()
if callable(track.OnSegmentFilter) and track.OnSegmentFilter(segment): if callable(track.OnSegmentFilter) and track.OnSegmentFilter(segment):
@ -430,7 +419,7 @@ class HLS:
finally: finally:
segment_key.put(newest_segment_key) segment_key.put(newest_segment_key)
if skip_event.is_set(): if DOWNLOAD_LICENCE_ONLY.is_set():
return -1 return -1
attempts = 1 attempts = 1
@ -456,9 +445,9 @@ class HLS:
segmented=True segmented=True
) )
break break
except Exception as ee: except Exception as e:
if stop_event.is_set() or attempts == 5: if DOWNLOAD_CANCELLED.is_set() or attempts == 5:
raise ee raise e
time.sleep(2) time.sleep(2)
attempts += 1 attempts += 1