1.4.4 upstream updates

This commit is contained in:
VineFeeder
2025-09-03 12:19:46 +01:00
parent 81d7ff4ba8
commit c44e076dec
26 changed files with 2040 additions and 239 deletions
+68 -41
View File
@@ -226,12 +226,6 @@ class dl:
default=False, default=False,
help="Skip downloading, only list available titles that would have been downloaded.", help="Skip downloading, only list available titles that would have been downloaded.",
) )
@click.option(
"--select-titles",
is_flag=True,
default=False,
help="Interactively select downloads from a list. Only use with Series to select Episodes",
)
@click.option( @click.option(
"--skip-dl", is_flag=True, default=False, help="Skip downloading while still retrieving the decryption keys." "--skip-dl", is_flag=True, default=False, help="Skip downloading while still retrieving the decryption keys."
) )
@@ -255,7 +249,17 @@ class dl:
) )
@click.option("--downloads", type=int, default=1, help="Amount of tracks to download concurrently.") @click.option("--downloads", type=int, default=1, help="Amount of tracks to download concurrently.")
@click.option("--no-cache", "no_cache", is_flag=True, default=False, help="Bypass title cache for this download.") @click.option("--no-cache", "no_cache", is_flag=True, default=False, help="Bypass title cache for this download.")
@click.option("--reset-cache", "reset_cache", is_flag=True, default=False, help="Clear title cache before fetching.") @click.option(
"--reset-cache", "reset_cache", is_flag=True, default=False, help="Clear title cache before fetching."
)
@click.option(
"--select-titles",
is_flag=True,
default=False,
help="Interactively select downloads from a list. Only use with Series to select Episodes",
)
@click.pass_context @click.pass_context
def cli(ctx: click.Context, **kwargs: Any) -> dl: def cli(ctx: click.Context, **kwargs: Any) -> dl:
return dl(ctx, **kwargs) return dl(ctx, **kwargs)
@@ -301,20 +305,8 @@ class dl:
if getattr(config, "downloader_map", None): if getattr(config, "downloader_map", None):
config.downloader = config.downloader_map.get(self.service, config.downloader) config.downloader = config.downloader_map.get(self.service, config.downloader)
with console.status("Loading DRM CDM...", spinner="dots"): if getattr(config, "decryption_map", None):
try: config.decryption = config.decryption_map.get(self.service, config.decryption)
self.cdm = self.get_cdm(self.service, self.profile)
except ValueError as e:
self.log.error(f"Failed to load CDM, {e}")
sys.exit(1)
if self.cdm:
if hasattr(self.cdm, "device_type") and self.cdm.device_type.name in ["ANDROID", "CHROME"]:
self.log.info(f"Loaded Widevine CDM: {self.cdm.system_id} (L{self.cdm.security_level})")
else:
self.log.info(
f"Loaded PlayReady CDM: {self.cdm.certificate_chain.get_name()} (L{self.cdm.security_level})"
)
with console.status("Loading Key Vaults...", spinner="dots"): with console.status("Loading Key Vaults...", spinner="dots"):
self.vaults = Vaults(self.service) self.vaults = Vaults(self.service)
@@ -354,6 +346,24 @@ class dl:
else: else:
self.log.debug("No vaults are currently active") self.log.debug("No vaults are currently active")
with console.status("Loading DRM CDM...", spinner="dots"):
try:
self.cdm = self.get_cdm(self.service, self.profile)
except ValueError as e:
self.log.error(f"Failed to load CDM, {e}")
sys.exit(1)
if self.cdm:
if isinstance(self.cdm, DecryptLabsRemoteCDM):
drm_type = "PlayReady" if self.cdm.is_playready else "Widevine"
self.log.info(f"Loaded {drm_type} Remote CDM: DecryptLabs (L{self.cdm.security_level})")
elif hasattr(self.cdm, "device_type") and self.cdm.device_type.name in ["ANDROID", "CHROME"]:
self.log.info(f"Loaded Widevine CDM: {self.cdm.system_id} (L{self.cdm.security_level})")
else:
self.log.info(
f"Loaded PlayReady CDM: {self.cdm.certificate_chain.get_name()} (L{self.cdm.security_level})"
)
self.proxy_providers = [] self.proxy_providers = []
if no_proxy: if no_proxy:
ctx.params["proxy"] = None ctx.params["proxy"] = None
@@ -441,7 +451,6 @@ class dl:
slow: bool, slow: bool,
list_: bool, list_: bool,
list_titles: bool, list_titles: bool,
select_titles: bool,
skip_dl: bool, skip_dl: bool,
export: Optional[Path], export: Optional[Path],
cdm_only: Optional[bool], cdm_only: Optional[bool],
@@ -450,6 +459,7 @@ class dl:
no_source: bool, no_source: bool,
workers: Optional[int], workers: Optional[int],
downloads: int, downloads: int,
select_titles: bool,
*_: Any, *_: Any,
**__: Any, **__: Any,
) -> None: ) -> None:
@@ -535,8 +545,6 @@ class dl:
del titles[i] del titles[i]
# end modification # end modification
for i, title in enumerate(titles): for i, title in enumerate(titles):
if isinstance(title, Episode) and wanted and f"{title.season}x{title.number}" not in wanted: if isinstance(title, Episode) and wanted and f"{title.season}x{title.number}" not in wanted:
continue continue
@@ -575,7 +583,7 @@ class dl:
else: else:
console.print(Padding("Search -> [bright_black]No match found[/]", (0, 5))) console.print(Padding("Search -> [bright_black]No match found[/]", (0, 5)))
if self.tmdb_id and getattr(self, 'search_source', None) != 'simkl': if self.tmdb_id and getattr(self, "search_source", None) != "simkl":
kind = "tv" if isinstance(title, Episode) else "movie" kind = "tv" if isinstance(title, Episode) else "movie"
tags.external_ids(self.tmdb_id, kind) tags.external_ids(self.tmdb_id, kind)
if self.tmdb_year: if self.tmdb_year:
@@ -913,7 +921,12 @@ class dl:
), ),
licence=partial( licence=partial(
service.get_playready_license service.get_playready_license
if isinstance(self.cdm, PlayReadyCdm) if (
isinstance(self.cdm, PlayReadyCdm)
or (
isinstance(self.cdm, DecryptLabsRemoteCDM) and self.cdm.is_playready
)
)
and hasattr(service, "get_playready_license") and hasattr(service, "get_playready_license")
else service.get_widevine_license, else service.get_widevine_license,
title=title, title=title,
@@ -1045,12 +1058,7 @@ class dl:
# Handle DRM decryption BEFORE repacking (must decrypt first!) # Handle DRM decryption BEFORE repacking (must decrypt first!)
service_name = service.__class__.__name__.upper() service_name = service.__class__.__name__.upper()
decryption_method = config.decryption_map.get(service_name, config.decryption) decryption_method = config.decryption_map.get(service_name, config.decryption)
use_mp4decrypt = decryption_method.lower() == "mp4decrypt" decrypt_tool = "mp4decrypt" if decryption_method.lower() == "mp4decrypt" else "Shaka Packager"
if use_mp4decrypt:
decrypt_tool = "mp4decrypt"
else:
decrypt_tool = "Shaka Packager"
drm_tracks = [track for track in title.tracks if track.drm] drm_tracks = [track for track in title.tracks if track.drm]
if drm_tracks: if drm_tracks:
@@ -1059,7 +1067,7 @@ class dl:
for track in drm_tracks: for track in drm_tracks:
drm = track.get_drm_for_cdm(self.cdm) drm = track.get_drm_for_cdm(self.cdm)
if drm and hasattr(drm, "decrypt"): if drm and hasattr(drm, "decrypt"):
drm.decrypt(track.path, use_mp4decrypt=use_mp4decrypt) drm.decrypt(track.path)
has_decrypted = True has_decrypted = True
events.emit(events.Types.TRACK_REPACKED, track=track) events.emit(events.Types.TRACK_REPACKED, track=track)
else: else:
@@ -1245,10 +1253,22 @@ class dl:
if not drm: if not drm:
return return
if isinstance(drm, Widevine) and not isinstance(self.cdm, WidevineCdm): if isinstance(drm, Widevine):
self.cdm = self.get_cdm(self.service, self.profile, drm="widevine") if not isinstance(self.cdm, (WidevineCdm, DecryptLabsRemoteCDM)) or (
elif isinstance(drm, PlayReady) and not isinstance(self.cdm, PlayReadyCdm): isinstance(self.cdm, DecryptLabsRemoteCDM) and self.cdm.is_playready
self.cdm = self.get_cdm(self.service, self.profile, drm="playready") ):
widevine_cdm = self.get_cdm(self.service, self.profile, drm="widevine")
if widevine_cdm:
self.log.info("Switching to Widevine CDM for Widevine content")
self.cdm = widevine_cdm
elif isinstance(drm, PlayReady):
if not isinstance(self.cdm, (PlayReadyCdm, DecryptLabsRemoteCDM)) or (
isinstance(self.cdm, DecryptLabsRemoteCDM) and not self.cdm.is_playready
):
playready_cdm = self.get_cdm(self.service, self.profile, drm="playready")
if playready_cdm:
self.log.info("Switching to PlayReady CDM for PlayReady content")
self.cdm = playready_cdm
if isinstance(drm, Widevine): if isinstance(drm, Widevine):
with self.DRM_TABLE_LOCK: with self.DRM_TABLE_LOCK:
@@ -1486,8 +1506,8 @@ class dl:
return Credential(*credentials) return Credential(*credentials)
return Credential.loads(credentials) # type: ignore return Credential.loads(credentials) # type: ignore
@staticmethod
def get_cdm( def get_cdm(
self,
service: str, service: str,
profile: Optional[str] = None, profile: Optional[str] = None,
drm: Optional[str] = None, drm: Optional[str] = None,
@@ -1521,10 +1541,18 @@ class dl:
cdm_api = next(iter(x for x in config.remote_cdm if x["name"] == cdm_name), None) cdm_api = next(iter(x for x in config.remote_cdm if x["name"] == cdm_name), None)
if cdm_api: if cdm_api:
is_decrypt_lab = True if cdm_api["type"] == "decrypt_labs" else False is_decrypt_lab = True if cdm_api.get("type") == "decrypt_labs" else False
if is_decrypt_lab:
del cdm_api["name"] del cdm_api["name"]
del cdm_api["type"] del cdm_api["type"]
return DecryptLabsRemoteCDM(service_name=service, **cdm_api) if is_decrypt_lab else RemoteCdm(**cdm_api)
# All DecryptLabs CDMs use DecryptLabsRemoteCDM
return DecryptLabsRemoteCDM(service_name=service, vaults=self.vaults, **cdm_api)
else:
del cdm_api["name"]
if "type" in cdm_api:
del cdm_api["type"]
return RemoteCdm(**cdm_api)
prd_path = config.directories.prds / f"{cdm_name}.prd" prd_path = config.directories.prds / f"{cdm_name}.prd"
if not prd_path.is_file(): if not prd_path.is_file():
@@ -1547,4 +1575,3 @@ class dl:
raise ValueError(f"{cdm_name}.wvd is an invalid or corrupt Widevine Device file, {e}") raise ValueError(f"{cdm_name}.wvd is an invalid or corrupt Widevine Device file, {e}")
return WidevineCdm.from_device(device) return WidevineCdm.from_device(device)
+1 -1
View File
@@ -1 +1 @@
__version__ = "1.4.2" __version__ = "1.4.4"
+29 -9
View File
@@ -10,13 +10,13 @@ from rich.padding import Padding
from rich.text import Text from rich.text import Text
from urllib3.exceptions import InsecureRequestWarning from urllib3.exceptions import InsecureRequestWarning
from . import __version__ from envied.core import __version__
from .commands import Commands from envied.core.commands import Commands
from .config import config from envied.core.config import config
from .console import ComfyRichHandler, console from envied.core.console import ComfyRichHandler, console
from .constants import context_settings from envied.core.constants import context_settings
from .update_checker import UpdateChecker from envied.core.update_checker import UpdateChecker
from .utilities import rotate_log_file from envied.core.utilities import rotate_log_file
LOGGING_PATH = None LOGGING_PATH = None
@@ -67,8 +67,8 @@ def main(version: bool, debug: bool, log_path: Path) -> None:
r"░▀▀▀░▀░▀░░▀░░▀▀▀░▀▀▀░▀▀░" + "\n" , r"░▀▀▀░▀░▀░░▀░░▀▀▀░▀▀▀░▀▀░" + "\n" ,
style="ascii.art", style="ascii.art",
), ),
"v 3.3.3 Copyright © 2019-2025 rlaphoenix"
+ f"\nv [repr.number]{__version__}[/] - unshackle", f"\nv [repr.number]{__version__}[/] - https://github.com/unshackle-dl/unshackle",
), ),
(1, 11, 1, 10), (1, 11, 1, 10),
expand=True, expand=True,
@@ -77,6 +77,26 @@ def main(version: bool, debug: bool, log_path: Path) -> None:
) )
'''if version:
return
if config.update_checks:
try:
latest_version = UpdateChecker.check_for_updates_sync(__version__)
if latest_version:
console.print(
f"\n[yellow]⚠️ Update available![/yellow] "
f"Current: {__version__} → Latest: [green]{latest_version}[/green]",
justify="center",
)
console.print(
"Visit: https://github.com/envied-dl/envied/releases/latest\n",
justify="center",
)
except Exception:
pass
'''
@atexit.register @atexit.register
def save_log(): def save_log():
if console.record and LOGGING_PATH: if console.record and LOGGING_PATH:
@@ -1,143 +1,658 @@
from __future__ import annotations
import base64 import base64
import secrets import secrets
from typing import Optional, Type, Union from typing import Any, Dict, List, Optional, Union
from uuid import UUID from uuid import UUID
import requests import requests
from pywidevine import PSSH, Device, DeviceTypes, Key, RemoteCdm from pywidevine.device import DeviceTypes
from pywidevine.license_protocol_pb2 import SignedDrmCertificate, SignedMessage from requests import Session
# Copyright 2024 by DevYukine. from envied.core.vaults import Vaults
class DecryptLabsRemoteCDM(RemoteCdm): class MockCertificateChain:
"""Mock certificate chain for PlayReady compatibility."""
def __init__(self, name: str):
self._name = name
def get_name(self) -> str:
return self._name
class Key:
"""Key object compatible with pywidevine."""
def __init__(self, kid: str, key: str, type_: str = "CONTENT"):
if isinstance(kid, str):
clean_kid = kid.replace("-", "")
if len(clean_kid) == 32:
self.kid = UUID(hex=clean_kid)
else:
self.kid = UUID(hex=clean_kid.ljust(32, "0"))
else:
self.kid = kid
if isinstance(key, str):
self.key = bytes.fromhex(key)
else:
self.key = key
self.type = type_
class DecryptLabsRemoteCDMExceptions:
"""Exception classes for compatibility with pywidevine CDM."""
class InvalidSession(Exception):
"""Raised when session ID is invalid."""
class TooManySessions(Exception):
"""Raised when session limit is reached."""
class InvalidInitData(Exception):
"""Raised when PSSH/init data is invalid."""
class InvalidLicenseType(Exception):
"""Raised when license type is invalid."""
class InvalidLicenseMessage(Exception):
"""Raised when license message is invalid."""
class InvalidContext(Exception):
"""Raised when session has no context data."""
class SignatureMismatch(Exception):
"""Raised when signature verification fails."""
class DecryptLabsRemoteCDM:
"""
Decrypt Labs Remote CDM implementation with intelligent caching system.
This class provides a drop-in replacement for pywidevine's local CDM using
Decrypt Labs' KeyXtractor API service, enhanced with smart caching logic
that minimizes unnecessary license requests.
Key Features:
- Compatible with both Widevine and PlayReady DRM schemes
- Intelligent caching that compares required vs. available keys
- Automatic key combination for mixed cache/license scenarios
- Seamless fallback to license requests when keys are missing
Intelligent Caching System:
1. DRM classes (PlayReady/Widevine) provide required KIDs via set_required_kids()
2. get_license_challenge() first checks for cached keys
3. If cached keys satisfy requirements, returns empty challenge (no license needed)
4. If keys are missing, makes targeted license request for remaining keys
5. parse_license() combines cached and license keys intelligently
"""
service_certificate_challenge = b"\x08\x04"
def __init__( def __init__(
self, self,
device_type: Union[DeviceTypes, str],
system_id: int,
security_level: int,
host: str,
secret: str, secret: str,
device_name: str, host: str = "https://keyxtractor.decryptlabs.com",
service_name: str, device_name: str = "ChromeCDM",
service_name: Optional[str] = None,
vaults: Optional[Vaults] = None,
device_type: Optional[str] = None,
system_id: Optional[int] = None,
security_level: Optional[int] = None,
**kwargs,
): ):
self.response_counter = 0 """
self.pssh = None Initialize Decrypt Labs Remote CDM for Widevine and PlayReady schemes.
self.api_session_ids = {}
self.license_request = None
self.service_name = service_name
self.keys = {}
try:
super().__init__(device_type, system_id, security_level, host, secret, device_name)
except Exception:
pass
self.req_session = requests.Session()
self.req_session.headers.update({"decrypt-labs-api-key": secret})
@classmethod Args:
def from_device(cls, device: Device) -> Type["DecryptLabsRemoteCDM"]: secret: Decrypt Labs API key (matches config format)
raise NotImplementedError("You cannot load a DecryptLabsRemoteCDM from a local Device file.") host: Decrypt Labs API host URL (matches config format)
device_name: DRM scheme (ChromeCDM, L1, L2 for Widevine; SL2, SL3 for PlayReady)
service_name: Service name for key caching and vault operations
vaults: Vaults instance for local key caching
device_type: Device type (CHROME, ANDROID, PLAYREADY) - for compatibility
system_id: System ID - for compatibility
security_level: Security level - for compatibility
"""
_ = kwargs
self.secret = secret
self.host = host.rstrip("/")
self.device_name = device_name
self.service_name = service_name or ""
self.vaults = vaults
self.uch = self.host != "https://keyxtractor.decryptlabs.com"
self._device_type_str = device_type
if device_type:
self.device_type = self._get_device_type_enum(device_type)
self._is_playready = (device_type and device_type.upper() == "PLAYREADY") or (device_name in ["SL2", "SL3"])
if self._is_playready:
self.system_id = system_id or 0
self.security_level = security_level or (2000 if device_name == "SL2" else 3000)
else:
self.system_id = system_id or 26830
self.security_level = security_level or 3
self._sessions: Dict[bytes, Dict[str, Any]] = {}
self._pssh_b64 = None
self._required_kids: Optional[List[str]] = None
self._http_session = Session()
self._http_session.headers.update(
{
"decrypt-labs-api-key": self.secret,
"Content-Type": "application/json",
"User-Agent": "envied-decrypt-labs-cdm/1.0",
}
)
def _get_device_type_enum(self, device_type: str):
"""Convert device type string to enum for compatibility."""
device_type_upper = device_type.upper()
if device_type_upper == "ANDROID":
return DeviceTypes.ANDROID
elif device_type_upper == "CHROME":
return DeviceTypes.CHROME
else:
return DeviceTypes.CHROME
@property
def is_playready(self) -> bool:
"""Check if this CDM is in PlayReady mode."""
return self._is_playready
@property
def certificate_chain(self) -> MockCertificateChain:
"""Mock certificate chain for PlayReady compatibility."""
return MockCertificateChain(f"{self.device_name}_Remote")
def set_pssh_b64(self, pssh_b64: str) -> None:
"""Store base64-encoded PSSH data for PlayReady compatibility."""
self._pssh_b64 = pssh_b64
def set_required_kids(self, kids: List[Union[str, UUID]]) -> None:
"""
Set the required Key IDs for intelligent caching decisions.
This method enables the CDM to make smart decisions about when to request
additional keys via license challenges. When cached keys are available,
the CDM will compare them against the required KIDs to determine if a
license request is still needed for missing keys.
Args:
kids: List of required Key IDs as UUIDs or hex strings
Note:
Should be called by DRM classes (PlayReady/Widevine) before making
license challenge requests to enable optimal caching behavior.
"""
self._required_kids = []
for kid in kids:
if isinstance(kid, UUID):
self._required_kids.append(str(kid).replace("-", "").lower())
else:
self._required_kids.append(str(kid).replace("-", "").lower())
def _generate_session_id(self) -> bytes:
"""Generate a unique session ID."""
return secrets.token_bytes(16)
def _get_init_data_from_pssh(self, pssh: Any) -> str:
"""Extract init data from various PSSH formats."""
if self.is_playready and self._pssh_b64:
return self._pssh_b64
if hasattr(pssh, "dumps"):
dumps_result = pssh.dumps()
if isinstance(dumps_result, str):
try:
base64.b64decode(dumps_result)
return dumps_result
except Exception:
return base64.b64encode(dumps_result.encode("utf-8")).decode("utf-8")
else:
return base64.b64encode(dumps_result).decode("utf-8")
elif hasattr(pssh, "raw"):
raw_data = pssh.raw
if isinstance(raw_data, str):
raw_data = raw_data.encode("utf-8")
return base64.b64encode(raw_data).decode("utf-8")
elif hasattr(pssh, "__class__") and "WrmHeader" in pssh.__class__.__name__:
if self.is_playready:
raise ValueError("PlayReady WRM header received but no PSSH B64 was set via set_pssh_b64()")
if hasattr(pssh, "raw_bytes"):
return base64.b64encode(pssh.raw_bytes).decode("utf-8")
elif hasattr(pssh, "bytes"):
return base64.b64encode(pssh.bytes).decode("utf-8")
else:
raise ValueError(f"Cannot extract PSSH data from WRM header type: {type(pssh)}")
else:
raise ValueError(f"Unsupported PSSH type: {type(pssh)}")
def open(self) -> bytes: def open(self) -> bytes:
# We stub this method to return a random session ID for now, later we save the api session id and resolve by our random generated one. """
return bytes.fromhex(secrets.token_hex(16)) Open a new CDM session.
Returns:
Session identifier as bytes
"""
session_id = self._generate_session_id()
self._sessions[session_id] = {
"service_certificate": None,
"keys": [],
"pssh": None,
"challenge": None,
"decrypt_labs_session_id": None,
}
return session_id
def close(self, session_id: bytes) -> None: def close(self, session_id: bytes) -> None:
# We stub this method to do nothing. """
pass Close a CDM session.
Args:
session_id: Session identifier
Raises:
ValueError: If session ID is invalid
"""
if session_id not in self._sessions:
raise DecryptLabsRemoteCDMExceptions.InvalidSession(f"Invalid session ID: {session_id.hex()}")
del self._sessions[session_id]
def get_service_certificate(self, session_id: bytes) -> Optional[bytes]:
"""
Get the service certificate for a session.
Args:
session_id: Session identifier
Returns:
Service certificate if set, None otherwise
Raises:
ValueError: If session ID is invalid
"""
if session_id not in self._sessions:
raise DecryptLabsRemoteCDMExceptions.InvalidSession(f"Invalid session ID: {session_id.hex()}")
return self._sessions[session_id]["service_certificate"]
def set_service_certificate(self, session_id: bytes, certificate: Optional[Union[bytes, str]]) -> str: def set_service_certificate(self, session_id: bytes, certificate: Optional[Union[bytes, str]]) -> str:
if isinstance(certificate, bytes): """
certificate = base64.b64encode(certificate).decode() Set the service certificate for a session.
# certificate needs to be base64 to be sent off to the API. Args:
# it needs to intentionally be kept as base64 encoded SignedMessage. session_id: Session identifier
certificate: Service certificate (bytes or base64 string)
self.req_session.signed_device_certificate = certificate Returns:
self.req_session.privacy_mode = True Certificate status message
return "success" Raises:
ValueError: If session ID is invalid
"""
if session_id not in self._sessions:
raise DecryptLabsRemoteCDMExceptions.InvalidSession(f"Invalid session ID: {session_id.hex()}")
def get_service_certificate(self, session_id: bytes) -> Optional[SignedDrmCertificate]: if certificate is None:
raise NotImplementedError("This method is not implemented in this CDM") self._sessions[session_id]["service_certificate"] = None
return "Removed"
if isinstance(certificate, str):
certificate = base64.b64decode(certificate)
self._sessions[session_id]["service_certificate"] = certificate
return "Successfully set Service Certificate"
def has_cached_keys(self, session_id: bytes) -> bool:
"""
Check if cached keys are available for the session.
Args:
session_id: Session identifier
Returns:
True if cached keys are available
Raises:
ValueError: If session ID is invalid
"""
if session_id not in self._sessions:
raise DecryptLabsRemoteCDMExceptions.InvalidSession(f"Invalid session ID: {session_id.hex()}")
session = self._sessions[session_id]
session_keys = session.get("keys", [])
return len(session_keys) > 0
def get_license_challenge( def get_license_challenge(
self, session_id: bytes, pssh: PSSH, license_type: str = "STREAMING", privacy_mode: bool = True self, session_id: bytes, pssh_or_wrm: Any, license_type: str = "STREAMING", privacy_mode: bool = True
) -> bytes: ) -> bytes:
self.pssh = pssh """
Generate a license challenge using Decrypt Labs API with intelligent caching.
res = self.session( This method implements smart caching logic that:
self.host + "/get-request", 1. First attempts to retrieve cached keys from the API
{ 2. If required KIDs are set, compares cached keys against requirements
"init_data": self.pssh.dumps(), 3. Only makes a license request if keys are missing
"service_certificate": self.req_session.signed_device_certificate, 4. Returns empty challenge if all required keys are cached
"scheme": "widevine",
"service": self.service_name,
},
)
self.license_request = res["challenge"] The intelligent caching works as follows:
self.api_session_ids[session_id] = res["session_id"] - With required KIDs set: Only requests license for missing keys
- Without required KIDs: Returns any available cached keys
- For PlayReady: Combines cached keys with license keys seamlessly
return base64.b64decode(self.license_request) Args:
session_id: Session identifier
pssh_or_wrm: PSSH object or WRM header (for PlayReady compatibility)
license_type: Type of license (STREAMING, OFFLINE, AUTOMATIC) - for compatibility only
privacy_mode: Whether to use privacy mode - for compatibility only
def parse_license(self, session_id: bytes, license_message: Union[SignedMessage, bytes, str]) -> None: Returns:
session_id_api = self.api_session_ids[session_id] License challenge as bytes, or empty bytes if cached keys satisfy requirements
if session_id not in self.keys:
self.keys[session_id] = []
session_keys = self.keys[session_id]
if isinstance(license_message, dict) and "keys" in license_message: Raises:
session_keys.extend( InvalidSession: If session ID is invalid
[ requests.RequestException: If API request fails
Key(kid=Key.kid_to_uuid(x["kid"]), type_=x.get("type", "CONTENT"), key=bytes.fromhex(x["key"]))
for x in license_message["keys"]
]
)
Note:
Call set_required_kids() before this method for optimal caching behavior.
"""
_ = license_type, privacy_mode
if session_id not in self._sessions:
raise DecryptLabsRemoteCDMExceptions.InvalidSession(f"Invalid session ID: {session_id.hex()}")
session = self._sessions[session_id]
session["pssh"] = pssh_or_wrm
init_data = self._get_init_data_from_pssh(pssh_or_wrm)
already_tried_cache = session.get("tried_cache", False)
request_data = {
"scheme": self.device_name,
"init_data": init_data,
"get_cached_keys_if_exists": not already_tried_cache,
}
if self.device_name in ["L1", "L2", "SL2", "SL3"] and self.service_name:
request_data["service"] = self.service_name
if session["service_certificate"]:
request_data["service_certificate"] = base64.b64encode(session["service_certificate"]).decode("utf-8")
response = self._http_session.post(f"{self.host}/get-request", json=request_data, timeout=30)
if response.status_code != 200:
raise requests.RequestException(f"API request failed: {response.status_code} {response.text}")
data = response.json()
if data.get("message") != "success":
error_msg = data.get("message", "Unknown error")
if "details" in data:
error_msg += f" - Details: {data['details']}"
if "error" in data:
error_msg += f" - Error: {data['error']}"
if "service_certificate is required" in str(data) and not session["service_certificate"]:
error_msg += " (No service certificate was provided to the CDM session)"
raise requests.RequestException(f"API error: {error_msg}")
message_type = data.get("message_type")
if message_type == "cached-keys" or "cached_keys" in data:
"""
Handle cached keys response from API.
When the API returns cached keys, we need to determine if they satisfy
our requirements or if we need to make an additional license request
for missing keys.
"""
cached_keys = data.get("cached_keys", [])
parsed_keys = self._parse_cached_keys(cached_keys)
session["keys"] = parsed_keys
session["tried_cache"] = True
if self._required_kids:
cached_kids = set()
for key in parsed_keys:
if isinstance(key, dict) and "kid" in key:
cached_kids.add(key["kid"].replace("-", "").lower())
required_kids = set(self._required_kids)
missing_kids = required_kids - cached_kids
if missing_kids:
session["cached_keys"] = parsed_keys
request_data["get_cached_keys_if_exists"] = False
response = self._http_session.post(f"{self.host}/get-request", json=request_data, timeout=30)
if response.status_code == 200:
data = response.json()
if data.get("message") == "success" and "challenge" in data:
challenge = base64.b64decode(data["challenge"])
session["challenge"] = challenge
session["decrypt_labs_session_id"] = data["session_id"]
return challenge
return b""
else: else:
res = self.session( return b""
self.host + "/decrypt-response",
{
"session_id": session_id_api,
"init_data": self.pssh.dumps(),
"license_request": self.license_request,
"license_response": license_message,
"scheme": "widevine",
},
)
original_keys = res["keys"].replace("\n", " ")
keys_separated = original_keys.split("--key ")
formatted_keys = []
for k in keys_separated:
if ":" in k:
key = k.strip()
formatted_keys.append(key)
for keys in formatted_keys:
session_keys.append(
(
Key(
kid=UUID(bytes=bytes.fromhex(keys.split(":")[0])),
type_="CONTENT",
key=bytes.fromhex(keys.split(":")[1]),
)
)
)
def get_keys(self, session_id: bytes, type_: Optional[Union[int, str]] = None) -> list[Key]:
return self.keys[session_id]
def session(self, url, data, retries=3):
res = self.req_session.post(url, json=data).json()
if res.get("message") != "success":
if "License Response Decryption Process Failed at the very beginning" in res.get("Error", ""):
if retries > 0:
return self.session(url, data, retries=retries - 1)
else: else:
raise ValueError(f"CDM API returned an error: {res['Error']}") return b""
else:
raise ValueError(f"CDM API returned an error: {res['Error']}")
return res if message_type == "license-request" or "challenge" in data:
challenge = base64.b64decode(data["challenge"])
session["challenge"] = challenge
session["decrypt_labs_session_id"] = data["session_id"]
return challenge
error_msg = f"Unexpected API response format. message_type={message_type}, available_fields={list(data.keys())}"
if data.get("message"):
error_msg = f"API response: {data['message']} - {error_msg}"
if "details" in data:
error_msg += f" - Details: {data['details']}"
if "error" in data:
error_msg += f" - Error: {data['error']}"
if already_tried_cache and data.get("message") == "success":
return b""
raise requests.RequestException(error_msg)
def parse_license(self, session_id: bytes, license_message: Union[bytes, str]) -> None:
"""
Parse license response using Decrypt Labs API with intelligent key combination.
For PlayReady content with partial cached keys, this method intelligently
combines the cached keys with newly obtained license keys, avoiding
duplicates while ensuring all required keys are available.
The key combination process:
1. Extracts keys from the license response
2. If cached keys exist (PlayReady), combines them with license keys
3. Removes duplicate keys by comparing normalized KIDs
4. Updates the session with the complete key set
Args:
session_id: Session identifier
license_message: License response from license server
Raises:
ValueError: If session ID is invalid or no challenge available
requests.RequestException: If API request fails
"""
if session_id not in self._sessions:
raise DecryptLabsRemoteCDMExceptions.InvalidSession(f"Invalid session ID: {session_id.hex()}")
session = self._sessions[session_id]
if session["keys"] and not (self.is_playready and "cached_keys" in session):
return
if not session.get("challenge") or not session.get("decrypt_labs_session_id"):
raise ValueError("No challenge available - call get_license_challenge first")
if isinstance(license_message, str):
if self.is_playready and license_message.strip().startswith("<?xml"):
license_message = license_message.encode("utf-8")
else:
try:
license_message = base64.b64decode(license_message)
except Exception:
license_message = license_message.encode("utf-8")
pssh = session["pssh"]
init_data = self._get_init_data_from_pssh(pssh)
license_request_b64 = base64.b64encode(session["challenge"]).decode("utf-8")
license_response_b64 = base64.b64encode(license_message).decode("utf-8")
request_data = {
"scheme": self.device_name,
"session_id": session["decrypt_labs_session_id"],
"init_data": init_data,
"license_request": license_request_b64,
"license_response": license_response_b64,
}
response = self._http_session.post(f"{self.host}/decrypt-response", json=request_data, timeout=30)
if response.status_code != 200:
raise requests.RequestException(f"License decrypt failed: {response.status_code} {response.text}")
data = response.json()
if data.get("message") != "success":
error_msg = data.get("message", "Unknown error")
if "error" in data:
error_msg += f" - Error: {data['error']}"
if "details" in data:
error_msg += f" - Details: {data['details']}"
raise requests.RequestException(f"License decrypt error: {error_msg}")
license_keys = self._parse_keys_response(data)
if self.is_playready and "cached_keys" in session:
"""
Combine cached keys with license keys for PlayReady content.
This ensures we have both the cached keys (obtained earlier) and
any additional keys from the license response, without duplicates.
"""
cached_keys = session.get("cached_keys", [])
all_keys = list(cached_keys)
for license_key in license_keys:
already_exists = False
license_kid = None
if isinstance(license_key, dict) and "kid" in license_key:
license_kid = license_key["kid"].replace("-", "").lower()
elif hasattr(license_key, "kid"):
license_kid = str(license_key.kid).replace("-", "").lower()
elif hasattr(license_key, "key_id"):
license_kid = str(license_key.key_id).replace("-", "").lower()
if license_kid:
for cached_key in cached_keys:
cached_kid = None
if isinstance(cached_key, dict) and "kid" in cached_key:
cached_kid = cached_key["kid"].replace("-", "").lower()
elif hasattr(cached_key, "kid"):
cached_kid = str(cached_key.kid).replace("-", "").lower()
elif hasattr(cached_key, "key_id"):
cached_kid = str(cached_key.key_id).replace("-", "").lower()
if cached_kid == license_kid:
already_exists = True
break
if not already_exists:
all_keys.append(license_key)
session["keys"] = all_keys
else:
session["keys"] = license_keys
if self.vaults and session["keys"]:
key_dict = {UUID(hex=key["kid"]): key["key"] for key in session["keys"] if key["type"] == "CONTENT"}
self.vaults.add_keys(key_dict)
def get_keys(self, session_id: bytes, type_: Optional[str] = None) -> List[Key]:
"""
Get keys from the session.
Args:
session_id: Session identifier
type_: Optional key type filter (CONTENT, SIGNING, etc.)
Returns:
List of Key objects
Raises:
InvalidSession: If session ID is invalid
"""
if session_id not in self._sessions:
raise DecryptLabsRemoteCDMExceptions.InvalidSession(f"Invalid session ID: {session_id.hex()}")
key_dicts = self._sessions[session_id]["keys"]
keys = [Key(kid=k["kid"], key=k["key"], type_=k["type"]) for k in key_dicts]
if type_:
keys = [key for key in keys if key.type == type_]
return keys
def _parse_cached_keys(self, cached_keys_data: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
"""Parse cached keys from API response.
Args:
cached_keys_data: List of cached key objects from API
Returns:
List of key dictionaries
"""
keys = []
try:
if cached_keys_data and isinstance(cached_keys_data, list):
for key_data in cached_keys_data:
if "kid" in key_data and "key" in key_data:
keys.append({"kid": key_data["kid"], "key": key_data["key"], "type": "CONTENT"})
except Exception:
pass
return keys
def _parse_keys_response(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
"""Parse keys from decrypt response."""
keys = []
if "keys" in data and isinstance(data["keys"], str):
keys_string = data["keys"]
for line in keys_string.split("\n"):
line = line.strip()
if line.startswith("--key "):
key_part = line[6:]
if ":" in key_part:
kid, key = key_part.split(":", 1)
keys.append({"kid": kid.strip(), "key": key.strip(), "type": "CONTENT"})
elif "keys" in data and isinstance(data["keys"], list):
for key_data in data["keys"]:
keys.append(
{"kid": key_data.get("kid"), "key": key_data.get("key"), "type": key_data.get("type", "CONTENT")}
)
return keys
__all__ = ["DecryptLabsRemoteCDM"]
+2 -2
View File
@@ -2,8 +2,8 @@ from typing import Optional
import click import click
from .config import config from envied.core.config import config
from .utilities import import_module_by_path from envied.core.utilities import import_module_by_path
_COMMANDS = sorted( _COMMANDS = sorted(
(path for path in config.directories.commands.glob("*.py") if path.stem.lower() != "__init__"), key=lambda x: x.stem (path for path in config.directories.commands.glob("*.py") if path.stem.lower() != "__init__"), key=lambda x: x.stem
@@ -93,7 +93,6 @@ class Config:
self.scene_naming: bool = kwargs.get("scene_naming", True) self.scene_naming: bool = kwargs.get("scene_naming", True)
self.series_year: bool = kwargs.get("series_year", True) self.series_year: bool = kwargs.get("series_year", True)
self.pssh_display: str = kwargs.get("pssh_display") or "fold" self.pssh_display: str = kwargs.get("pssh_display") or "fold"
self.title_cache_time: int = kwargs.get("title_cache_time", 1800) # 30 minutes default self.title_cache_time: int = kwargs.get("title_cache_time", 1800) # 30 minutes default
self.title_cache_max_retention: int = kwargs.get("title_cache_max_retention", 86400) # 24 hours default self.title_cache_max_retention: int = kwargs.get("title_cache_max_retention", 86400) # 24 hours default
self.title_cache_enabled: bool = kwargs.get("title_cache_enabled", True) self.title_cache_enabled: bool = kwargs.get("title_cache_enabled", True)
@@ -224,15 +224,51 @@ class PlayReady:
def kids(self) -> list[UUID]: def kids(self) -> list[UUID]:
return self._kids return self._kids
def get_content_keys(self, cdm: PlayReadyCdm, certificate: Callable, licence: Callable) -> None: def _extract_keys_from_cdm(self, cdm: PlayReadyCdm, session_id: bytes) -> dict:
for kid in self.kids: """Extract keys from CDM session with cross-library compatibility.
if kid in self.content_keys:
Args:
cdm: CDM instance
session_id: Session identifier
Returns:
Dictionary mapping KID UUIDs to hex keys
"""
keys = {}
for key in cdm.get_keys(session_id):
if hasattr(key, "key_id"):
kid = key.key_id
elif hasattr(key, "kid"):
kid = key.kid
else:
continue continue
if hasattr(key, "key") and hasattr(key.key, "hex"):
key_hex = key.key.hex()
elif hasattr(key, "key") and isinstance(key.key, bytes):
key_hex = key.key.hex()
elif hasattr(key, "key") and isinstance(key.key, str):
key_hex = key.key
else:
continue
keys[kid] = key_hex
return keys
def get_content_keys(self, cdm: PlayReadyCdm, certificate: Callable, licence: Callable) -> None:
session_id = cdm.open() session_id = cdm.open()
try: try:
challenge = cdm.get_license_challenge(session_id, self.pssh.wrm_headers[0]) if hasattr(cdm, "set_pssh_b64") and self.pssh_b64:
license_res = licence(challenge=challenge) cdm.set_pssh_b64(self.pssh_b64)
if hasattr(cdm, "set_required_kids"):
cdm.set_required_kids(self.kids)
challenge = cdm.get_license_challenge(session_id, self.pssh.wrm_headers[0])
if challenge:
try:
license_res = licence(challenge=challenge)
if isinstance(license_res, bytes): if isinstance(license_res, bytes):
license_str = license_res.decode(errors="ignore") license_str = license_res.decode(errors="ignore")
else: else:
@@ -245,7 +281,10 @@ class PlayReady:
pass pass
cdm.parse_license(session_id, license_str) cdm.parse_license(session_id, license_str)
keys = {key.key_id: key.key.hex() for key in cdm.get_keys(session_id)} except Exception:
raise
keys = self._extract_keys_from_cdm(cdm, session_id)
self.content_keys.update(keys) self.content_keys.update(keys)
finally: finally:
cdm.close(session_id) cdm.close(session_id)
@@ -253,12 +292,11 @@ class PlayReady:
if not self.content_keys: if not self.content_keys:
raise PlayReady.Exceptions.EmptyLicense("No Content Keys were within the License") raise PlayReady.Exceptions.EmptyLicense("No Content Keys were within the License")
def decrypt(self, path: Path, use_mp4decrypt: bool = False) -> None: def decrypt(self, path: Path) -> None:
""" """
Decrypt a Track with PlayReady DRM. Decrypt a Track with PlayReady DRM.
Args: Args:
path: Path to the encrypted file to decrypt path: Path to the encrypted file to decrypt
use_mp4decrypt: If True, use mp4decrypt instead of Shaka Packager
Raises: Raises:
EnvironmentError if the required decryption executable could not be found. EnvironmentError if the required decryption executable could not be found.
ValueError if the track has not yet been downloaded. ValueError if the track has not yet been downloaded.
@@ -270,7 +308,9 @@ class PlayReady:
if not path or not path.exists(): if not path or not path.exists():
raise ValueError("Tried to decrypt a file that does not exist.") raise ValueError("Tried to decrypt a file that does not exist.")
if use_mp4decrypt: decrypter = str(getattr(config, "decryption", "")).lower()
if decrypter == "mp4decrypt":
return self._decrypt_with_mp4decrypt(path) return self._decrypt_with_mp4decrypt(path)
else: else:
return self._decrypt_with_shaka_packager(path) return self._decrypt_with_shaka_packager(path)
@@ -185,7 +185,15 @@ class Widevine:
if cert and hasattr(cdm, "set_service_certificate"): if cert and hasattr(cdm, "set_service_certificate"):
cdm.set_service_certificate(session_id, cert) cdm.set_service_certificate(session_id, cert)
cdm.parse_license(session_id, licence(challenge=cdm.get_license_challenge(session_id, self.pssh))) if hasattr(cdm, "set_required_kids"):
cdm.set_required_kids(self.kids)
challenge = cdm.get_license_challenge(session_id, self.pssh)
if hasattr(cdm, "has_cached_keys") and cdm.has_cached_keys(session_id):
pass
else:
cdm.parse_license(session_id, licence(challenge=challenge))
self.content_keys = {key.kid: key.key.hex() for key in cdm.get_keys(session_id, "CONTENT")} self.content_keys = {key.kid: key.key.hex() for key in cdm.get_keys(session_id, "CONTENT")}
if not self.content_keys: if not self.content_keys:
@@ -213,9 +221,17 @@ class Widevine:
if cert and hasattr(cdm, "set_service_certificate"): if cert and hasattr(cdm, "set_service_certificate"):
cdm.set_service_certificate(session_id, cert) cdm.set_service_certificate(session_id, cert)
if hasattr(cdm, "set_required_kids"):
cdm.set_required_kids(self.kids)
challenge = cdm.get_license_challenge(session_id, self.pssh)
if hasattr(cdm, "has_cached_keys") and cdm.has_cached_keys(session_id):
pass
else:
cdm.parse_license( cdm.parse_license(
session_id, session_id,
licence(session_id=session_id, challenge=cdm.get_license_challenge(session_id, self.pssh)), licence(session_id=session_id, challenge=challenge),
) )
self.content_keys = {key.kid: key.key.hex() for key in cdm.get_keys(session_id, "CONTENT")} self.content_keys = {key.kid: key.key.hex() for key in cdm.get_keys(session_id, "CONTENT")}
@@ -227,12 +243,11 @@ class Widevine:
finally: finally:
cdm.close(session_id) cdm.close(session_id)
def decrypt(self, path: Path, use_mp4decrypt: bool = False) -> None: def decrypt(self, path: Path) -> None:
""" """
Decrypt a Track with Widevine DRM. Decrypt a Track with Widevine DRM.
Args: Args:
path: Path to the encrypted file to decrypt path: Path to the encrypted file to decrypt
use_mp4decrypt: If True, use mp4decrypt instead of Shaka Packager
Raises: Raises:
EnvironmentError if the required decryption executable could not be found. EnvironmentError if the required decryption executable could not be found.
ValueError if the track has not yet been downloaded. ValueError if the track has not yet been downloaded.
@@ -244,7 +259,9 @@ class Widevine:
if not path or not path.exists(): if not path or not path.exists():
raise ValueError("Tried to decrypt a file that does not exist.") raise ValueError("Tried to decrypt a file that does not exist.")
if use_mp4decrypt: decrypter = str(getattr(config, "decryption", "")).lower()
if decrypter == "mp4decrypt":
return self._decrypt_with_mp4decrypt(path) return self._decrypt_with_mp4decrypt(path)
else: else:
return self._decrypt_with_shaka_packager(path) return self._decrypt_with_shaka_packager(path)
@@ -420,6 +420,15 @@ class Track:
for drm in self.drm: for drm in self.drm:
if isinstance(drm, PlayReady): if isinstance(drm, PlayReady):
return drm return drm
elif hasattr(cdm, 'is_playready'):
if cdm.is_playready:
for drm in self.drm:
if isinstance(drm, PlayReady):
return drm
else:
for drm in self.drm:
if isinstance(drm, Widevine):
return drm
return self.drm[0] return self.drm[0]
+3 -3
View File
@@ -22,9 +22,9 @@ from langcodes import Language, closest_match
from pymp4.parser import Box from pymp4.parser import Box
from unidecode import unidecode from unidecode import unidecode
from .cacher import Cacher from envied.core.cacher import Cacher
from .config import config from envied.core.config import config
from .constants import LANGUAGE_MAX_DISTANCE from envied.core.constants import LANGUAGE_MAX_DISTANCE
def rotate_log_file(log_path: Path, keep: int = 20) -> Path: def rotate_log_file(log_path: Path, keep: int = 20) -> Path:
+25 -10
View File
@@ -8,6 +8,7 @@ import tempfile
from difflib import SequenceMatcher from difflib import SequenceMatcher
from pathlib import Path from pathlib import Path
from typing import Optional, Tuple from typing import Optional, Tuple
from xml.sax.saxutils import escape
import requests import requests
from requests.adapters import HTTPAdapter, Retry from requests.adapters import HTTPAdapter, Retry
@@ -289,9 +290,9 @@ def _apply_tags(path: Path, tags: dict[str, str]) -> None:
log.debug("mkvpropedit not found on PATH; skipping tags") log.debug("mkvpropedit not found on PATH; skipping tags")
return return
log.debug("Applying tags to %s: %s", path, tags) log.debug("Applying tags to %s: %s", path, tags)
xml_lines = ["<?xml version='1.0' encoding='UTF-8'?>", "<Tags>", " <Tag>", " <Targets/>"] xml_lines = ['<?xml version="1.0" encoding="UTF-8"?>', "<Tags>", " <Tag>", " <Targets/>"]
for name, value in tags.items(): for name, value in tags.items():
xml_lines.append(f" <Simple><Name>{name}</Name><String>{value}</String></Simple>") xml_lines.append(f" <Simple><Name>{escape(name)}</Name><String>{escape(value)}</String></Simple>")
xml_lines.extend([" </Tag>", "</Tags>"]) xml_lines.extend([" </Tag>", "</Tags>"])
with tempfile.NamedTemporaryFile("w", suffix=".xml", delete=False) as f: with tempfile.NamedTemporaryFile("w", suffix=".xml", delete=False) as f:
f.write("\n".join(xml_lines)) f.write("\n".join(xml_lines))
@@ -349,13 +350,25 @@ def tag_file(path: Path, title: Title, tmdb_id: Optional[int] | None = None) ->
if simkl_tmdb_id: if simkl_tmdb_id:
tmdb_id = simkl_tmdb_id tmdb_id = simkl_tmdb_id
# Handle TV show data from Simkl
if simkl_data.get("type") == "episode" and "show" in simkl_data:
show_ids = simkl_data.get("show", {}).get("ids", {}) show_ids = simkl_data.get("show", {}).get("ids", {})
if show_ids.get("imdb"): if show_ids.get("imdb"):
standard_tags["IMDB"] = f"https://www.imdb.com/title/{show_ids['imdb']}" standard_tags["IMDB"] = show_ids["imdb"]
if show_ids.get("tvdb"): if show_ids.get("tvdb"):
standard_tags["TVDB"] = f"https://thetvdb.com/dereferrer/series/{show_ids['tvdb']}" standard_tags["TVDB2"] = f"series/{show_ids['tvdb']}"
if show_ids.get("tmdbtv"): if show_ids.get("tmdbtv"):
standard_tags["TMDB"] = f"https://www.themoviedb.org/tv/{show_ids['tmdbtv']}" standard_tags["TMDB"] = f"tv/{show_ids['tmdbtv']}"
# Handle movie data from Simkl
elif simkl_data.get("type") == "movie" and "movie" in simkl_data:
movie_ids = simkl_data.get("movie", {}).get("ids", {})
if movie_ids.get("imdb"):
standard_tags["IMDB"] = movie_ids["imdb"]
if movie_ids.get("tvdb"):
standard_tags["TVDB2"] = f"movies/{movie_ids['tvdb']}"
if movie_ids.get("tmdb"):
standard_tags["TMDB"] = f"movie/{movie_ids['tmdb']}"
# Use TMDB API for additional metadata (either from provided ID or Simkl lookup) # Use TMDB API for additional metadata (either from provided ID or Simkl lookup)
api_key = _api_key() api_key = _api_key()
@@ -373,8 +386,8 @@ def tag_file(path: Path, title: Title, tmdb_id: Optional[int] | None = None) ->
_apply_tags(path, custom_tags) _apply_tags(path, custom_tags)
return return
tmdb_url = f"https://www.themoviedb.org/{'movie' if kind == 'movie' else 'tv'}/{tmdb_id}" prefix = "movie" if kind == "movie" else "tv"
standard_tags["TMDB"] = tmdb_url standard_tags["TMDB"] = f"{prefix}/{tmdb_id}"
try: try:
ids = external_ids(tmdb_id, kind) ids = external_ids(tmdb_id, kind)
except requests.RequestException as exc: except requests.RequestException as exc:
@@ -385,11 +398,13 @@ def tag_file(path: Path, title: Title, tmdb_id: Optional[int] | None = None) ->
imdb_id = ids.get("imdb_id") imdb_id = ids.get("imdb_id")
if imdb_id: if imdb_id:
standard_tags["IMDB"] = f"https://www.imdb.com/title/{imdb_id}" standard_tags["IMDB"] = imdb_id
tvdb_id = ids.get("tvdb_id") tvdb_id = ids.get("tvdb_id")
if tvdb_id: if tvdb_id:
tvdb_prefix = "movies" if kind == "movie" else "series" if kind == "movie":
standard_tags["TVDB"] = f"https://thetvdb.com/dereferrer/{tvdb_prefix}/{tvdb_id}" standard_tags["TVDB2"] = f"movies/{tvdb_id}"
else:
standard_tags["TVDB2"] = f"series/{tvdb_id}"
merged_tags = { merged_tags = {
**custom_tags, **custom_tags,
@@ -26,7 +26,7 @@ from .schemes import EntityAuthenticationSchemes # noqa: F401
from .schemes import KeyExchangeSchemes from .schemes import KeyExchangeSchemes
from .schemes.EntityAuthentication import EntityAuthentication from .schemes.EntityAuthentication import EntityAuthentication
from .schemes.KeyExchangeRequest import KeyExchangeRequest from .schemes.KeyExchangeRequest import KeyExchangeRequest
# from vinetrimmer.utils.widevine.device import RemoteDevice # from vinetrimmer.utils.wienvied.device import RemoteDevice
class MSL: class MSL:
log = logging.getLogger("MSL") log = logging.getLogger("MSL")
@@ -45,7 +45,7 @@ class RTE(Service):
""" """
# GEOFENCE = ("ie",) #GEOFENCE = ("ie",)
@staticmethod @staticmethod
@click.command(name="RTE", short_help="https://www.rte.ie/player/", help=__doc__) @click.command(name="RTE", short_help="https://www.rte.ie/player/", help=__doc__)
@@ -4,4 +4,4 @@ headers:
endpoints: endpoints:
base_url: https://www.rte.ie base_url: https://www.rte.ie
feed: https://feed.entertainment.tv.theplatform.eu feed: https://feed.entertainment.tv.theplatform.eu
license: https://widevine.entitlement.eu.theplatform.com/wv/web/ModularDrm license: https://wienvied.entitlement.eu.theplatform.com/wv/web/ModularDrm
@@ -166,11 +166,11 @@ class STV(Service):
source source
for source in data["sources"] for source in data["sources"]
if source.get("type") == "application/dash+xml" if source.get("type") == "application/dash+xml"
and source.get("key_systems").get("com.widevine.alpha")), and source.get("key_systems").get("com.wienvied.alpha")),
None, None,
) )
self.license = key_systems["key_systems"]["com.widevine.alpha"]["license_url"] if key_systems else None self.license = key_systems["key_systems"]["com.wienvied.alpha"]["license_url"] if key_systems else None
manifest = self.trim_duration(source_manifest) manifest = self.trim_duration(source_manifest)
tracks = DASH.from_text(manifest, source_manifest).to_tracks(title.language) tracks = DASH.from_text(manifest, source_manifest).to_tracks(title.language)
@@ -20,7 +20,7 @@ def get_widevine_license_url(manifest_str):
try: try:
data = json.loads(manifest_str) data = json.loads(manifest_str)
for source in data.get("sources", []): for source in data.get("sources", []):
widevine = source.get("key_systems", {}).get("com.widevine.alpha") widevine = source.get("key_systems", {}).get("com.wienvied.alpha")
if widevine and "license_url" in widevine: if widevine and "license_url" in widevine:
return widevine["license_url"] return widevine["license_url"]
except json.JSONDecodeError: except json.JSONDecodeError:
@@ -172,14 +172,14 @@ class TVNZ(Service):
) )
self.license = next(( self.license = next((
x["key_systems"]["com.widevine.alpha"]["license_url"] x["key_systems"]["com.wienvied.alpha"]["license_url"]
for x in data["sources"] for x in data["sources"]
if x.get("key_systems").get("com.widevine.alpha")), if x.get("key_systems").get("com.wienvied.alpha")),
None, None,
) )
source_manifest = next(( source_manifest = next((
x["src"] for x in data["sources"] x["src"] for x in data["sources"]
if x.get("key_systems").get("com.widevine.alpha")), if x.get("key_systems").get("com.wienvied.alpha")),
None, None,
) )
@@ -120,14 +120,14 @@ class UKTV(Service):
data = r.json() data = r.json()
self.license = next(( self.license = next((
x["key_systems"]["com.widevine.alpha"]["license_url"] x["key_systems"]["com.wienvied.alpha"]["license_url"]
for x in data["sources"] for x in data["sources"]
if x.get("key_systems").get("com.widevine.alpha")), if x.get("key_systems").get("com.wienvied.alpha")),
None, None,
) )
source_manifest = next(( source_manifest = next((
x["src"] for x in data["sources"] x["src"] for x in data["sources"]
if x.get("key_systems").get("com.widevine.alpha")), if x.get("key_systems").get("com.wienvied.alpha")),
None, None,
) )
if not self.license or not source_manifest: if not self.license or not source_manifest:
+184
View File
@@ -0,0 +1,184 @@
from typing import Iterator, Optional, Union
from uuid import UUID
from requests import Session
from envied.core import __version__
from envied.core.vault import Vault
class API(Vault):
"""Key Vault using a simple RESTful HTTP API call."""
def __init__(self, name: str, uri: str, token: str, no_push: bool = False):
super().__init__(name, no_push)
self.uri = uri.rstrip("/")
self.session = Session()
self.session.headers.update({"User-Agent": f"envied v{__version__}"})
self.session.headers.update({"Authorization": f"Bearer {token}"})
def get_key(self, kid: Union[UUID, str], service: str) -> Optional[str]:
if isinstance(kid, UUID):
kid = kid.hex
data = self.session.get(
url=f"{self.uri}/{service.lower()}/{kid}", headers={"Accept": "application/json"}
).json()
code = int(data.get("code", 0))
message = data.get("message")
error = {
0: None,
1: Exceptions.AuthRejected,
2: Exceptions.TooManyRequests,
3: Exceptions.ServiceTagInvalid,
4: Exceptions.KeyIdInvalid,
}.get(code, ValueError)
if error:
raise error(f"{message} ({code})")
content_key = data.get("content_key")
if not content_key:
return None
if not isinstance(content_key, str):
raise ValueError(f"Expected {content_key} to be {str}, was {type(content_key)}")
return content_key
def get_keys(self, service: str) -> Iterator[tuple[str, str]]:
page = 1
while True:
data = self.session.get(
url=f"{self.uri}/{service.lower()}",
params={"page": page, "total": 10},
headers={"Accept": "application/json"},
).json()
code = int(data.get("code", 0))
message = data.get("message")
error = {
0: None,
1: Exceptions.AuthRejected,
2: Exceptions.TooManyRequests,
3: Exceptions.PageInvalid,
4: Exceptions.ServiceTagInvalid,
}.get(code, ValueError)
if error:
raise error(f"{message} ({code})")
content_keys = data.get("content_keys")
if content_keys:
if not isinstance(content_keys, dict):
raise ValueError(f"Expected {content_keys} to be {dict}, was {type(content_keys)}")
for key_id, key in content_keys.items():
yield key_id, key
pages = int(data["pages"])
if pages <= page:
break
page += 1
def add_key(self, service: str, kid: Union[UUID, str], key: str) -> bool:
if isinstance(kid, UUID):
kid = kid.hex
data = self.session.post(
url=f"{self.uri}/{service.lower()}/{kid}", json={"content_key": key}, headers={"Accept": "application/json"}
).json()
code = int(data.get("code", 0))
message = data.get("message")
error = {
0: None,
1: Exceptions.AuthRejected,
2: Exceptions.TooManyRequests,
3: Exceptions.ServiceTagInvalid,
4: Exceptions.KeyIdInvalid,
5: Exceptions.ContentKeyInvalid,
}.get(code, ValueError)
if error:
raise error(f"{message} ({code})")
# the kid:key was new to the vault (optional)
added = bool(data.get("added"))
# the key for kid was changed/updated (optional)
updated = bool(data.get("updated"))
return added or updated
def add_keys(self, service: str, kid_keys: dict[Union[UUID, str], str]) -> int:
data = self.session.post(
url=f"{self.uri}/{service.lower()}",
json={"content_keys": {str(kid).replace("-", ""): key for kid, key in kid_keys.items()}},
headers={"Accept": "application/json"},
).json()
code = int(data.get("code", 0))
message = data.get("message")
error = {
0: None,
1: Exceptions.AuthRejected,
2: Exceptions.TooManyRequests,
3: Exceptions.ServiceTagInvalid,
4: Exceptions.KeyIdInvalid,
5: Exceptions.ContentKeyInvalid,
}.get(code, ValueError)
if error:
raise error(f"{message} ({code})")
# each kid:key that was new to the vault (optional)
added = int(data.get("added"))
# each key for a kid that was changed/updated (optional)
updated = int(data.get("updated"))
return added + updated
def get_services(self) -> Iterator[str]:
data = self.session.post(url=self.uri, headers={"Accept": "application/json"}).json()
code = int(data.get("code", 0))
message = data.get("message")
error = {
0: None,
1: Exceptions.AuthRejected,
2: Exceptions.TooManyRequests,
}.get(code, ValueError)
if error:
raise error(f"{message} ({code})")
service_list = data.get("service_list", [])
if not isinstance(service_list, list):
raise ValueError(f"Expected {service_list} to be {list}, was {type(service_list)}")
for service in service_list:
yield service
class Exceptions:
class AuthRejected(Exception):
"""Authentication Error Occurred, is your token valid? Do you have permission to make this call?"""
class TooManyRequests(Exception):
"""Rate Limited; Sent too many requests in a given amount of time."""
class PageInvalid(Exception):
"""Requested page does not exist."""
class ServiceTagInvalid(Exception):
"""The Service Tag is invalid."""
class KeyIdInvalid(Exception):
"""The Key ID is invalid."""
class ContentKeyInvalid(Exception):
"""The Content Key is invalid."""
+337
View File
@@ -0,0 +1,337 @@
import json
from enum import Enum
from typing import Iterator, Optional, Union
from uuid import UUID
from requests import Session
from envied.core import __version__
from envied.core.vault import Vault
class InsertResult(Enum):
FAILURE = 0
SUCCESS = 1
ALREADY_EXISTS = 2
class HTTP(Vault):
"""Key Vault using HTTP API with support for both query parameters and JSON payloads."""
def __init__(
self,
name: str,
host: str,
password: str,
username: Optional[str] = None,
api_mode: str = "query",
no_push: bool = False,
):
"""
Initialize HTTP Vault.
Args:
name: Vault name
host: Host URL
password: Password for query mode or API token for json mode
username: Username (required for query mode, ignored for json mode)
api_mode: "query" for query parameters or "json" for JSON API
no_push: If True, this vault will not receive pushed keys
"""
super().__init__(name, no_push)
self.url = host
self.password = password
self.username = username
self.api_mode = api_mode.lower()
self.current_title = None
self.session = Session()
self.session.headers.update({"User-Agent": f"envied v{__version__}"})
self.api_session_id = None
# Validate configuration based on mode
if self.api_mode == "query" and not self.username:
raise ValueError("Username is required for query mode")
elif self.api_mode not in ["query", "json"]:
raise ValueError("api_mode must be either 'query' or 'json'")
def request(self, method: str, params: dict = None) -> dict:
"""Make a request to the JSON API vault."""
if self.api_mode != "json":
raise ValueError("request method is only available in json mode")
request_payload = {
"method": method,
"params": {
**(params or {}),
"session_id": self.api_session_id,
},
"token": self.password,
}
r = self.session.post(self.url, json=request_payload)
if r.status_code == 404:
return {"status": "not_found"}
if not r.ok:
raise ValueError(f"API returned HTTP Error {r.status_code}: {r.reason.title()}")
try:
res = r.json()
except json.JSONDecodeError:
if r.status_code == 404:
return {"status": "not_found"}
raise ValueError(f"API returned an invalid response: {r.text}")
if res.get("status_code") != 200:
raise ValueError(f"API returned an error: {res['status_code']} - {res['message']}")
if session_id := res.get("message", {}).get("session_id"):
self.api_session_id = session_id
return res.get("message", res)
def get_key(self, kid: Union[UUID, str], service: str) -> Optional[str]:
if isinstance(kid, UUID):
kid = kid.hex
if self.api_mode == "json":
try:
params = {
"kid": kid,
"service": service.lower(),
}
response = self.request("GetKey", params)
if response.get("status") == "not_found":
return None
keys = response.get("keys", [])
for key_entry in keys:
if isinstance(key_entry, str) and ":" in key_entry:
entry_kid, entry_key = key_entry.split(":", 1)
if entry_kid == kid:
return entry_key
elif isinstance(key_entry, dict):
if key_entry.get("kid") == kid:
return key_entry.get("key")
except Exception as e:
print(f"Failed to get key ({e.__class__.__name__}: {e})")
return None
return None
else: # query mode
response = self.session.get(
self.url,
params={"service": service.lower(), "username": self.username, "password": self.password, "kid": kid},
)
data = response.json()
if data.get("status_code") != 200 or not data.get("keys"):
return None
return data["keys"][0]["key"]
def get_keys(self, service: str) -> Iterator[tuple[str, str]]:
if self.api_mode == "json":
# JSON API doesn't support getting all keys, so return empty iterator
# This will cause the copy command to rely on the API's internal duplicate handling
return iter([])
else: # query mode
response = self.session.get(
self.url, params={"service": service.lower(), "username": self.username, "password": self.password}
)
data = response.json()
if data.get("status_code") != 200 or not data.get("keys"):
return
for key_entry in data["keys"]:
yield key_entry["kid"], key_entry["key"]
def add_key(self, service: str, kid: Union[UUID, str], key: str) -> bool:
if not key or key.count("0") == len(key):
raise ValueError("You cannot add a NULL Content Key to a Vault.")
if isinstance(kid, UUID):
kid = kid.hex
title = getattr(self, "current_title", None)
if self.api_mode == "json":
try:
response = self.request(
"InsertKey",
{
"kid": kid,
"key": key,
"service": service.lower(),
"title": title,
},
)
if response.get("status") == "not_found":
return False
return response.get("inserted", False)
except Exception:
return False
else: # query mode
response = self.session.get(
self.url,
params={
"service": service.lower(),
"username": self.username,
"password": self.password,
"kid": kid,
"key": key,
"title": title,
},
)
data = response.json()
return data.get("status_code") == 200
def add_keys(self, service: str, kid_keys: dict[Union[UUID, str], str]) -> int:
for kid, key in kid_keys.items():
if not key or key.count("0") == len(key):
raise ValueError("You cannot add a NULL Content Key to a Vault.")
processed_kid_keys = {
str(kid).replace("-", "") if isinstance(kid, UUID) else kid: key for kid, key in kid_keys.items()
}
inserted_count = 0
title = getattr(self, "current_title", None)
if self.api_mode == "json":
for kid, key in processed_kid_keys.items():
try:
response = self.request(
"InsertKey",
{
"kid": kid,
"key": key,
"service": service.lower(),
"title": title,
},
)
if response.get("status") == "not_found":
continue
if response.get("inserted", False):
inserted_count += 1
except Exception:
continue
else: # query mode
for kid, key in processed_kid_keys.items():
response = self.session.get(
self.url,
params={
"service": service.lower(),
"username": self.username,
"password": self.password,
"kid": kid,
"key": key,
"title": title,
},
)
data = response.json()
if data.get("status_code") == 200 and data.get("inserted", True):
inserted_count += 1
return inserted_count
def get_services(self) -> Iterator[str]:
if self.api_mode == "json":
try:
response = self.request("GetServices")
services = response.get("services", [])
for service in services:
yield service
except Exception:
return iter([])
else: # query mode
response = self.session.get(
self.url, params={"username": self.username, "password": self.password, "list_services": True}
)
data = response.json()
if data.get("status_code") != 200:
return
services = data.get("services", [])
for service in services:
yield service
def set_title(self, title: str):
"""
Set a title to be used for the next key insertions.
This is optional and will be sent with add_key requests if available.
"""
self.current_title = title
def insert_key_with_result(
self, service: str, kid: Union[UUID, str], key: str, title: Optional[str] = None
) -> InsertResult:
"""
Insert a key and return detailed result information.
This method provides more granular feedback than the standard add_key method.
Available in both API modes.
"""
if not key or key.count("0") == len(key):
raise ValueError("You cannot add a NULL Content Key to a Vault.")
if isinstance(kid, UUID):
kid = kid.hex
if title is None:
title = getattr(self, "current_title", None)
if self.api_mode == "json":
try:
response = self.request(
"InsertKey",
{
"kid": kid,
"key": key,
"service": service.lower(),
"title": title,
},
)
if response.get("status") == "not_found":
return InsertResult.FAILURE
if response.get("inserted", False):
return InsertResult.SUCCESS
else:
return InsertResult.ALREADY_EXISTS
except Exception:
return InsertResult.FAILURE
else: # query mode
response = self.session.get(
self.url,
params={
"service": service.lower(),
"username": self.username,
"password": self.password,
"kid": kid,
"key": key,
"title": title,
},
)
try:
data = response.json()
if data.get("status_code") == 200:
if data.get("inserted", True):
return InsertResult.SUCCESS
else:
return InsertResult.ALREADY_EXISTS
else:
return InsertResult.FAILURE
except Exception:
return InsertResult.FAILURE
@@ -0,0 +1,215 @@
import json
from typing import Iterator, Optional, Union
from uuid import UUID
from enum import Enum
from requests import Session
from envied.core import __version__
from envied.core.vault import Vault
class InsertResult(Enum):
FAILURE = 0
SUCCESS = 1
ALREADY_EXISTS = 2
class HTTPAPI(Vault):
"""Key Vault using a structured HTTP API with JSON payloads and token authentication."""
def __init__(self, name: str, host: str, password: str):
super().__init__(name)
self.url = host
self.password = password # This is the API token
self.current_title = None
self.session = Session()
self.session.headers.update({"User-Agent": f"Devine v{__version__}"})
self.api_session_id = None
def request(self, method: str, params: dict = None) -> dict:
"""Make a request to the HTTPAPI vault."""
request_payload = {
"method": method,
"params": {
**(params or {}),
"session_id": self.api_session_id,
},
"token": self.password,
}
r = self.session.post(self.url, json=request_payload)
if r.status_code == 404:
return {"status": "not_found"}
if not r.ok:
raise ValueError(f"API returned HTTP Error {r.status_code}: {r.reason.title()}")
try:
res = r.json()
except json.JSONDecodeError:
if r.status_code == 404:
return {"status": "not_found"}
raise ValueError(f"API returned an invalid response: {r.text}")
if res.get("status_code") != 200:
raise ValueError(f"API returned an error: {res['status_code']} - {res['message']}")
if session_id := res.get("message", {}).get("session_id"):
self.api_session_id = session_id
return res.get("message", res)
def get_key(self, kid: Union[UUID, str], service: str) -> Optional[str]:
if isinstance(kid, UUID):
kid = kid.hex
try:
title = getattr(self, "current_title", None)
response = self.request(
"GetKey",
{
"kid": kid,
"service": service.lower(),
"title": title,
},
)
if response.get("status") == "not_found":
return None
keys = response.get("keys", [])
for key_entry in keys:
if key_entry["kid"] == kid:
return key_entry["key"]
except Exception as e:
print(f"Failed to get key ({e.__class__.__name__}: {e})")
return None
return None
def get_keys(self, service: str) -> Iterator[tuple[str, str]]:
"""Get all keys for a service - this may need to be implemented based on your API."""
try:
response = self.request(
"GetKeys",
{
"service": service.lower(),
},
)
keys = response.get("keys", [])
for key_entry in keys:
yield key_entry["kid"], key_entry["key"]
except Exception:
return iter([])
def add_key(self, service: str, kid: Union[UUID, str], key: str) -> bool:
if not key or key.count("0") == len(key):
raise ValueError("You cannot add a NULL Content Key to a Vault.")
if isinstance(kid, UUID):
kid = kid.hex
title = getattr(self, "current_title", None)
try:
response = self.request(
"InsertKey",
{
"kid": kid,
"key": key,
"service": service.lower(),
"title": title,
},
)
if response.get("status") == "not_found":
return False
return response.get("inserted", False)
except Exception:
return False
def add_keys(self, service: str, kid_keys: dict[Union[UUID, str], str]) -> int:
for kid, key in kid_keys.items():
if not key or key.count("0") == len(key):
raise ValueError("You cannot add a NULL Content Key to a Vault.")
processed_kid_keys = {
str(kid).replace("-", "") if isinstance(kid, UUID) else kid: key for kid, key in kid_keys.items()
}
inserted_count = 0
title = getattr(self, "current_title", None)
for kid, key in processed_kid_keys.items():
try:
response = self.request(
"InsertKey",
{
"kid": kid,
"key": key,
"service": service.lower(),
"title": title,
},
)
if response.get("status") == "not_found":
continue
if response.get("inserted", False):
inserted_count += 1
except Exception:
continue
return inserted_count
def get_services(self) -> Iterator[str]:
"""Get available services - this may need to be implemented based on your API."""
try:
response = self.request("GetServices")
services = response.get("services", [])
for service in services:
yield service
except Exception:
return iter([])
def set_title(self, title: str):
"""
Set a title to be used for the next key insertions.
This is optional and will be sent with add_key requests if available.
"""
self.current_title = title
def insert_key_with_result(
self, service: str, kid: Union[UUID, str], key: str, title: Optional[str] = None
) -> InsertResult:
"""
Insert a key and return detailed result information.
This method provides more granular feedback than the standard add_key method.
"""
if not key or key.count("0") == len(key):
raise ValueError("You cannot add a NULL Content Key to a Vault.")
if isinstance(kid, UUID):
kid = kid.hex
if title is None:
title = getattr(self, "current_title", None)
try:
response = self.request(
"InsertKey",
{
"kid": kid,
"key": key,
"service": service.lower(),
"title": title,
},
)
if response.get("status") == "not_found":
return InsertResult.FAILURE
if response.get("inserted", False):
return InsertResult.SUCCESS
else:
return InsertResult.ALREADY_EXISTS
except Exception:
return InsertResult.FAILURE
+244
View File
@@ -0,0 +1,244 @@
import threading
from typing import Iterator, Optional, Union
from uuid import UUID
import pymysql
from pymysql.cursors import DictCursor
from envied.core.services import Services
from envied.core.vault import Vault
class MySQL(Vault):
"""Key Vault using a remotely-accessed mysql database connection."""
def __init__(self, name: str, host: str, database: str, username: str, no_push: bool = False, **kwargs):
"""
All extra arguments provided via **kwargs will be sent to pymysql.connect.
This can be used to provide more specific connection information.
"""
super().__init__(name, no_push)
self.slug = f"{host}:{database}:{username}"
self.conn_factory = ConnectionFactory(
dict(host=host, db=database, user=username, cursorclass=DictCursor, **kwargs)
)
self.permissions = self.get_permissions()
if not self.has_permission("SELECT"):
raise PermissionError(f"MySQL vault {self.slug} has no SELECT permission.")
def get_key(self, kid: Union[UUID, str], service: str) -> Optional[str]:
if not self.has_table(service):
# no table, no key, simple
return None
if isinstance(kid, UUID):
kid = kid.hex
conn = self.conn_factory.get()
cursor = conn.cursor()
try:
cursor.execute(
# TODO: SQL injection risk
f"SELECT `id`, `key_` FROM `{service}` WHERE `kid`=%s AND `key_`!=%s",
(kid, "0" * 32),
)
cek = cursor.fetchone()
if not cek:
return None
return cek["key_"]
finally:
cursor.close()
def get_keys(self, service: str) -> Iterator[tuple[str, str]]:
if not self.has_table(service):
# no table, no keys, simple
return None
conn = self.conn_factory.get()
cursor = conn.cursor()
try:
cursor.execute(
# TODO: SQL injection risk
f"SELECT `kid`, `key_` FROM `{service}` WHERE `key_`!=%s",
("0" * 32,),
)
for row in cursor.fetchall():
yield row["kid"], row["key_"]
finally:
cursor.close()
def add_key(self, service: str, kid: Union[UUID, str], key: str) -> bool:
if not key or key.count("0") == len(key):
raise ValueError("You cannot add a NULL Content Key to a Vault.")
if not self.has_permission("INSERT", table=service):
raise PermissionError(f"MySQL vault {self.slug} has no INSERT permission.")
if not self.has_table(service):
try:
self.create_table(service)
except PermissionError:
return False
if isinstance(kid, UUID):
kid = kid.hex
conn = self.conn_factory.get()
cursor = conn.cursor()
try:
cursor.execute(
# TODO: SQL injection risk
f"SELECT `id` FROM `{service}` WHERE `kid`=%s AND `key_`=%s",
(kid, key),
)
if cursor.fetchone():
# table already has this exact KID:KEY stored
return True
cursor.execute(
# TODO: SQL injection risk
f"INSERT INTO `{service}` (kid, key_) VALUES (%s, %s)",
(kid, key),
)
finally:
conn.commit()
cursor.close()
return True
def add_keys(self, service: str, kid_keys: dict[Union[UUID, str], str]) -> int:
for kid, key in kid_keys.items():
if not key or key.count("0") == len(key):
raise ValueError("You cannot add a NULL Content Key to a Vault.")
if not self.has_permission("INSERT", table=service):
raise PermissionError(f"MySQL vault {self.slug} has no INSERT permission.")
if not self.has_table(service):
try:
self.create_table(service)
except PermissionError:
return 0
if not isinstance(kid_keys, dict):
raise ValueError(f"The kid_keys provided is not a dictionary, {kid_keys!r}")
if not all(isinstance(kid, (str, UUID)) and isinstance(key_, str) for kid, key_ in kid_keys.items()):
raise ValueError("Expecting dict with Key of str/UUID and value of str.")
if any(isinstance(kid, UUID) for kid, key_ in kid_keys.items()):
kid_keys = {kid.hex if isinstance(kid, UUID) else kid: key_ for kid, key_ in kid_keys.items()}
conn = self.conn_factory.get()
cursor = conn.cursor()
try:
cursor.executemany(
# TODO: SQL injection risk
f"INSERT IGNORE INTO `{service}` (kid, key_) VALUES (%s, %s)",
kid_keys.items(),
)
return cursor.rowcount
finally:
conn.commit()
cursor.close()
def get_services(self) -> Iterator[str]:
conn = self.conn_factory.get()
cursor = conn.cursor()
try:
cursor.execute("SHOW TABLES")
for table in cursor.fetchall():
# each entry has a key named `Tables_in_<db name>`
yield Services.get_tag(list(table.values())[0])
finally:
cursor.close()
def has_table(self, name: str) -> bool:
"""Check if the Vault has a Table with the specified name."""
conn = self.conn_factory.get()
cursor = conn.cursor()
try:
cursor.execute(
"SELECT count(TABLE_NAME) FROM information_schema.TABLES WHERE TABLE_SCHEMA=%s AND TABLE_NAME=%s",
(conn.db, name),
)
return list(cursor.fetchone().values())[0] == 1
finally:
cursor.close()
def create_table(self, name: str):
"""Create a Table with the specified name if not yet created."""
if self.has_table(name):
return
if not self.has_permission("CREATE"):
raise PermissionError(f"MySQL vault {self.slug} has no CREATE permission.")
conn = self.conn_factory.get()
cursor = conn.cursor()
try:
cursor.execute(
# TODO: SQL injection risk
f"""
CREATE TABLE IF NOT EXISTS {name} (
id int AUTO_INCREMENT PRIMARY KEY,
kid VARCHAR(64) NOT NULL,
key_ VARCHAR(64) NOT NULL,
UNIQUE(kid, key_)
);
"""
)
finally:
conn.commit()
cursor.close()
def get_permissions(self) -> list:
"""Get and parse Grants to a more easily usable list tuple array."""
conn = self.conn_factory.get()
cursor = conn.cursor()
try:
cursor.execute("SHOW GRANTS")
grants = cursor.fetchall()
grants = [next(iter(x.values())) for x in grants]
grants = [tuple(x[6:].split(" TO ")[0].split(" ON ")) for x in list(grants)]
grants = [
(
list(map(str.strip, perms.replace("ALL PRIVILEGES", "*").split(","))),
location.replace("`", "").split("."),
)
for perms, location in grants
]
return grants
finally:
conn.commit()
cursor.close()
def has_permission(self, operation: str, database: Optional[str] = None, table: Optional[str] = None) -> bool:
"""Check if the current connection has a specific permission."""
grants = [x for x in self.permissions if x[0] == ["*"] or operation.upper() in x[0]]
if grants and database:
grants = [x for x in grants if x[1][0] in (database, "*")]
if grants and table:
grants = [x for x in grants if x[1][1] in (table, "*")]
return bool(grants)
class ConnectionFactory:
def __init__(self, con: dict):
self._con = con
self._store = threading.local()
def _create_connection(self) -> pymysql.Connection:
return pymysql.connect(**self._con)
def get(self) -> pymysql.Connection:
if not hasattr(self._store, "conn"):
self._store.conn = self._create_connection()
return self._store.conn
+179
View File
@@ -0,0 +1,179 @@
import sqlite3
import threading
from pathlib import Path
from sqlite3 import Connection
from typing import Iterator, Optional, Union
from uuid import UUID
from envied.core.services import Services
from envied.core.vault import Vault
class SQLite(Vault):
"""Key Vault using a locally-accessed sqlite DB file."""
def __init__(self, name: str, path: Union[str, Path], no_push: bool = False):
super().__init__(name, no_push)
self.path = Path(path).expanduser()
# TODO: Use a DictCursor or such to get fetches as dict?
self.conn_factory = ConnectionFactory(self.path)
def get_key(self, kid: Union[UUID, str], service: str) -> Optional[str]:
if not self.has_table(service):
# no table, no key, simple
return None
if isinstance(kid, UUID):
kid = kid.hex
conn = self.conn_factory.get()
cursor = conn.cursor()
try:
cursor.execute(f"SELECT `id`, `key_` FROM `{service}` WHERE `kid`=? AND `key_`!=?", (kid, "0" * 32))
cek = cursor.fetchone()
if not cek:
return None
return cek[1]
finally:
cursor.close()
def get_keys(self, service: str) -> Iterator[tuple[str, str]]:
if not self.has_table(service):
# no table, no keys, simple
return None
conn = self.conn_factory.get()
cursor = conn.cursor()
try:
cursor.execute(f"SELECT `kid`, `key_` FROM `{service}` WHERE `key_`!=?", ("0" * 32,))
for kid, key_ in cursor.fetchall():
yield kid, key_
finally:
cursor.close()
def add_key(self, service: str, kid: Union[UUID, str], key: str) -> bool:
if not key or key.count("0") == len(key):
raise ValueError("You cannot add a NULL Content Key to a Vault.")
if not self.has_table(service):
self.create_table(service)
if isinstance(kid, UUID):
kid = kid.hex
conn = self.conn_factory.get()
cursor = conn.cursor()
try:
cursor.execute(
# TODO: SQL injection risk
f"SELECT `id` FROM `{service}` WHERE `kid`=? AND `key_`=?",
(kid, key),
)
if cursor.fetchone():
# table already has this exact KID:KEY stored
return True
cursor.execute(
# TODO: SQL injection risk
f"INSERT INTO `{service}` (kid, key_) VALUES (?, ?)",
(kid, key),
)
finally:
conn.commit()
cursor.close()
return True
def add_keys(self, service: str, kid_keys: dict[Union[UUID, str], str]) -> int:
for kid, key in kid_keys.items():
if not key or key.count("0") == len(key):
raise ValueError("You cannot add a NULL Content Key to a Vault.")
if not self.has_table(service):
self.create_table(service)
if not isinstance(kid_keys, dict):
raise ValueError(f"The kid_keys provided is not a dictionary, {kid_keys!r}")
if not all(isinstance(kid, (str, UUID)) and isinstance(key_, str) for kid, key_ in kid_keys.items()):
raise ValueError("Expecting dict with Key of str/UUID and value of str.")
if any(isinstance(kid, UUID) for kid, key_ in kid_keys.items()):
kid_keys = {kid.hex if isinstance(kid, UUID) else kid: key_ for kid, key_ in kid_keys.items()}
conn = self.conn_factory.get()
cursor = conn.cursor()
try:
cursor.executemany(
# TODO: SQL injection risk
f"INSERT OR IGNORE INTO `{service}` (kid, key_) VALUES (?, ?)",
kid_keys.items(),
)
return cursor.rowcount
finally:
conn.commit()
cursor.close()
def get_services(self) -> Iterator[str]:
conn = self.conn_factory.get()
cursor = conn.cursor()
try:
cursor.execute("SELECT name FROM sqlite_master WHERE type='table';")
for (name,) in cursor.fetchall():
if name != "sqlite_sequence":
yield Services.get_tag(name)
finally:
cursor.close()
def has_table(self, name: str) -> bool:
"""Check if the Vault has a Table with the specified name."""
conn = self.conn_factory.get()
cursor = conn.cursor()
try:
cursor.execute("SELECT count(name) FROM sqlite_master WHERE type='table' AND name=?", (name,))
return cursor.fetchone()[0] == 1
finally:
cursor.close()
def create_table(self, name: str):
"""Create a Table with the specified name if not yet created."""
if self.has_table(name):
return
conn = self.conn_factory.get()
cursor = conn.cursor()
try:
cursor.execute(
# TODO: SQL injection risk
f"""
CREATE TABLE IF NOT EXISTS {name} (
"id" INTEGER NOT NULL UNIQUE,
"kid" TEXT NOT NULL COLLATE NOCASE,
"key_" TEXT NOT NULL COLLATE NOCASE,
PRIMARY KEY("id" AUTOINCREMENT),
UNIQUE("kid", "key_")
);
"""
)
finally:
conn.commit()
cursor.close()
class ConnectionFactory:
def __init__(self, path: Union[str, Path]):
self._path = path
self._store = threading.local()
def _create_connection(self) -> Connection:
return sqlite3.connect(self._path)
def get(self) -> Connection:
if not hasattr(self._store, "conn"):
self._store.conn = self._create_connection()
return self._store.conn
@@ -1 +1 @@
__version__ = "1.0.1" __version__ = "1.0.2"
@@ -406,7 +406,7 @@ class BaseLoader:
if self.TERMINAL_RESET and not self.BATCH_DOWNLOAD: if self.TERMINAL_RESET and not self.BATCH_DOWNLOAD:
console.print(f"[{catppuccin_mocha['text2']}]Preparing to reset Terminal[/]") console.print(f"[{catppuccin_mocha['text2']}]Preparing to reset Terminal[/]")
time.sleep(5) time.sleep(5)
self.reset_terminal() #self.reset_terminal()
console.print(f"[{catppuccin_mocha['pink']}]Ready![/]") console.print(f"[{catppuccin_mocha['pink']}]Ready![/]")