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

Pass save path to DRM decrypt functions directly 

This is required in segmented scenarios when multi-threaded where the same `track.path` would be get and set from possibly at the same time. It's also just better logically to do it this way.
This commit is contained in:
rlaphoenix 2023-02-21 16:09:35 +00:00
parent 8268825ba8
commit 314079c75f
5 changed files with 32 additions and 34 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
devine/commands/dl.py
View File

@ -630,7 +630,6 @@ class dl:
service.session.headers, service.session.headers,
proxy if track.needs_proxy else None proxy if track.needs_proxy else None
)) ))
track.path = save_path
if not track.drm and isinstance(track, (Video, Audio)): if not track.drm and isinstance(track, (Video, Audio)):
try: try:
@ -644,7 +643,9 @@ class dl:
if isinstance(drm, Widevine): if isinstance(drm, Widevine):
# license and grab content keys # license and grab content keys
prepare_drm(drm) prepare_drm(drm)
drm.decrypt(track) drm.decrypt(save_path)
track.drm = None
track.path = save_path
if callable(track.OnDecrypted): if callable(track.OnDecrypted):
track.OnDecrypted(track) track.OnDecrypted(track)
else: else:

18
devine/core/drm/clearkey.py
View File

@ -1,5 +1,7 @@
from __future__ import annotations from __future__ import annotations
import shutil
from pathlib import Path
from typing import Optional, Union from typing import Optional, Union
from urllib.parse import urljoin from urllib.parse import urljoin
@ -7,8 +9,6 @@ import requests
from Cryptodome.Cipher import AES from Cryptodome.Cipher import AES
from m3u8.model import Key from m3u8.model import Key
from devine.core.constants import TrackT
class ClearKey: class ClearKey:
"""AES Clear Key DRM System.""" """AES Clear Key DRM System."""
@ -34,20 +34,20 @@ class ClearKey:
self.key: bytes = key self.key: bytes = key
self.iv: bytes = iv self.iv: bytes = iv
def decrypt(self, track: TrackT) -> None: def decrypt(self, path: Path) -> None:
"""Decrypt a Track with AES Clear Key DRM.""" """Decrypt a Track with AES Clear Key DRM."""
if not track.path or not track.path.exists(): if not path or not path.exists():
raise ValueError("Tried to decrypt a track that has not yet been downloaded.") raise ValueError("Tried to decrypt a file that does not exist.")
decrypted = AES. \ decrypted = AES. \
new(self.key, AES.MODE_CBC, self.iv). \ new(self.key, AES.MODE_CBC, self.iv). \
decrypt(track.path.read_bytes()) decrypt(path.read_bytes())
decrypted_path = track.path.with_suffix(f".decrypted{track.path.suffix}") decrypted_path = path.with_suffix(f".decrypted{path.suffix}")
decrypted_path.write_bytes(decrypted) decrypted_path.write_bytes(decrypted)
track.swap(decrypted_path) path.unlink()
track.drm = None shutil.move(decrypted_path, path)
@classmethod @classmethod
def from_m3u_key(cls, m3u_key: Key, proxy: Optional[str] = None) -> ClearKey: def from_m3u_key(cls, m3u_key: Key, proxy: Optional[str] = None) -> ClearKey:

19
devine/core/drm/widevine.py
View File

@ -1,8 +1,10 @@
from __future__ import annotations from __future__ import annotations
import base64 import base64
import shutil
import subprocess import subprocess
import sys import sys
from pathlib import Path
from typing import Any, Callable, Optional, Union from typing import Any, Callable, Optional, Union
from uuid import UUID from uuid import UUID
@ -14,7 +16,7 @@ from pywidevine.pssh import PSSH
from requests import Session from requests import Session
from devine.core.config import config from devine.core.config import config
from devine.core.constants import AnyTrack, TrackT from devine.core.constants import AnyTrack
from devine.core.utilities import get_binary_path, get_boxes from devine.core.utilities import get_binary_path, get_boxes
from devine.core.utils.subprocess import ffprobe from devine.core.utils.subprocess import ffprobe
@ -212,7 +214,7 @@ class Widevine:
finally: finally:
cdm.close(session_id) cdm.close(session_id)
def decrypt(self, track: TrackT) -> None: def decrypt(self, path: Path) -> None:
""" """
Decrypt a Track with Widevine DRM. Decrypt a Track with Widevine DRM.
Raises: Raises:
@ -227,15 +229,15 @@ class Widevine:
executable = get_binary_path("shaka-packager", f"packager-{platform}", f"packager-{platform}-x64") executable = get_binary_path("shaka-packager", f"packager-{platform}", f"packager-{platform}-x64")
if not executable: if not executable:
raise EnvironmentError("Shaka Packager executable not found but is required.") raise EnvironmentError("Shaka Packager executable not found but is required.")
if not track.path or not track.path.exists(): if not path or not path.exists():
raise ValueError("Tried to decrypt a track that has not yet been downloaded.") raise ValueError("Tried to decrypt a file that does not exist.")
decrypted_path = track.path.with_suffix(f".decrypted{track.path.suffix}") decrypted_path = path.with_suffix(f".decrypted{path.suffix}")
config.directories.temp.mkdir(parents=True, exist_ok=True) config.directories.temp.mkdir(parents=True, exist_ok=True)
try: try:
subprocess.check_call([ subprocess.check_call([
executable, executable,
f"input={track.path},stream=0,output={decrypted_path}", f"input={path},stream=0,output={decrypted_path}",
"--enable_raw_key_decryption", "--keys", "--enable_raw_key_decryption", "--keys",
",".join([ ",".join([
*[ *[
@ -252,8 +254,9 @@ class Widevine:
], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) ], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
except subprocess.CalledProcessError as e: except subprocess.CalledProcessError as e:
raise subprocess.SubprocessError(f"Failed to Decrypt! Shaka Packager Error: {e}") raise subprocess.SubprocessError(f"Failed to Decrypt! Shaka Packager Error: {e}")
track.swap(decrypted_path)
track.drm = None path.unlink()
shutil.move(decrypted_path, path)
class Exceptions: class Exceptions:
class PSSHNotFound(Exception): class PSSHNotFound(Exception):

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

@ -414,12 +414,9 @@ class DASH:
session.headers, session.headers,
proxy proxy
)) ))
# TODO: More like `segment.path`, but this will do for now
# Needed for the drm.decrypt() call couple lines down
track.path = segment_save_path
if isinstance(track, Audio) or init_data: if isinstance(track, Audio) or init_data:
with open(track.path, "rb+") as f: with open(segment_save_path, "rb+") as f:
segment_data = f.read() segment_data = f.read()
if isinstance(track, Audio): if isinstance(track, Audio):
# fix audio decryption on ATVP by fixing the sample description index # fix audio decryption on ATVP by fixing the sample description index
@ -437,7 +434,8 @@ class DASH:
if drm: if drm:
# TODO: What if the manifest does not mention DRM, but has DRM # TODO: What if the manifest does not mention DRM, but has DRM
drm.decrypt(track) drm.decrypt(segment_save_path)
track.drm = None
if callable(track.OnDecrypted): if callable(track.OnDecrypted):
track.OnDecrypted(track) track.OnDecrypted(track)
elif segment_list is not None: elif segment_list is not None:
@ -485,12 +483,9 @@ class DASH:
proxy, proxy,
byte_range=segment_url.get("mediaRange") byte_range=segment_url.get("mediaRange")
)) ))
# TODO: More like `segment.path`, but this will do for now
# Needed for the drm.decrypt() call couple lines down
track.path = segment_save_path
if isinstance(track, Audio) or init_data: if isinstance(track, Audio) or init_data:
with open(track.path, "rb+") as f: with open(segment_save_path, "rb+") as f:
segment_data = f.read() segment_data = f.read()
if isinstance(track, Audio): if isinstance(track, Audio):
# fix audio decryption on ATVP by fixing the sample description index # fix audio decryption on ATVP by fixing the sample description index
@ -508,7 +503,8 @@ class DASH:
if drm: if drm:
# TODO: What if the manifest does not mention DRM, but has DRM # TODO: What if the manifest does not mention DRM, but has DRM
drm.decrypt(track) drm.decrypt(segment_save_path)
track.drm = None
if callable(track.OnDecrypted): if callable(track.OnDecrypted):
track.OnDecrypted(track) track.OnDecrypted(track)
elif segment_base is not None or base_url: elif segment_base is not None or base_url:

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

@ -269,12 +269,9 @@ class HLS:
session.headers, session.headers,
proxy proxy
)) ))
# TODO: More like `segment.path`, but this will do for now
# Needed for the drm.decrypt() call couple lines down
track.path = segment_save_path
if isinstance(track, Audio) or init_data: if isinstance(track, Audio) or init_data:
with open(track.path, "rb+") as f: with open(segment_save_path, "rb+") as f:
segment_data = f.read() segment_data = f.read()
if isinstance(track, Audio): if isinstance(track, Audio):
# fix audio decryption on ATVP by fixing the sample description index # fix audio decryption on ATVP by fixing the sample description index
@ -291,7 +288,8 @@ class HLS:
f.write(segment_data) f.write(segment_data)
if last_segment_key[0]: if last_segment_key[0]:
last_segment_key[0].decrypt(track) last_segment_key[0].decrypt(segment_save_path)
track.drm = None
if callable(track.OnDecrypted): if callable(track.OnDecrypted):
track.OnDecrypted(track) track.OnDecrypted(track)