From 79cdbc007cc817a0a7282c2f61dbfd887b7b3896 Mon Sep 17 00:00:00 2001 From: rlaphoenix Date: Wed, 8 Nov 2023 22:38:32 +0000 Subject: [PATCH] Remove `Types` shortcut from Device, rename to DeviceTypes This is because a static linter cannot recognize a class variable as a type. If we instead directly reference the enum, it can. --- pywidevine/cdm.py | 12 ++++++------ pywidevine/device.py | 13 ++++++------- pywidevine/main.py | 6 +++--- pywidevine/remotecdm.py | 10 +++++----- 4 files changed, 20 insertions(+), 21 deletions(-) diff --git a/pywidevine/cdm.py b/pywidevine/cdm.py index 11b8aa0..db83f45 100644 --- a/pywidevine/cdm.py +++ b/pywidevine/cdm.py @@ -18,7 +18,7 @@ from Crypto.Signature import pss from Crypto.Util import Padding from google.protobuf.message import DecodeError -from pywidevine.device import Device +from pywidevine.device import Device, DeviceTypes from pywidevine.exceptions import (InvalidContext, InvalidInitData, InvalidLicenseMessage, InvalidLicenseType, InvalidSession, NoKeysLoaded, SignatureMismatch, TooManySessions) from pywidevine.key import Key @@ -66,7 +66,7 @@ class Cdm: def __init__( self, - device_type: Union[Device.Types, str], + device_type: Union[DeviceTypes, str], system_id: int, security_level: int, client_id: ClientIdentification, @@ -76,9 +76,9 @@ class Cdm: if not device_type: raise ValueError("Device Type must be provided") if isinstance(device_type, str): - device_type = Device.Types[device_type] - if not isinstance(device_type, Device.Types): - raise TypeError(f"Expected device_type to be a {Device.Types!r} not {device_type!r}") + device_type = DeviceTypes[device_type] + if not isinstance(device_type, DeviceTypes): + raise TypeError(f"Expected device_type to be a {DeviceTypes!r} not {device_type!r}") if not system_id: raise ValueError("System ID must be provided") @@ -306,7 +306,7 @@ class Cdm: f"Available values: {LicenseType.keys()}" ) - if self.device_type == Device.Types.ANDROID: + if self.device_type == DeviceTypes.ANDROID: # OEMCrypto's request_id seems to be in AES CTR Counter block form with no suffix # Bytes 5-8 does not seem random, in real tests they have been consecutive \x00 or \xFF # Real example: A0DCE548000000000500000000000000 diff --git a/pywidevine/device.py b/pywidevine/device.py index 1deeaf4..4625cf5 100644 --- a/pywidevine/device.py +++ b/pywidevine/device.py @@ -17,7 +17,7 @@ from google.protobuf.message import DecodeError from pywidevine.license_protocol_pb2 import ClientIdentification, DrmCertificate, FileHashes, SignedDrmCertificate -class _Types(Enum): +class DeviceTypes(Enum): CHROME = 1 ANDROID = 2 @@ -36,7 +36,7 @@ class _Structures: "version" / Const(Int8ub, 2), "type_" / CEnum( Int8ub, - **{t.name: t.value for t in _Types} + **{t.name: t.value for t in DeviceTypes} ), "security_level" / Int8ub, "flags" / Padded(1, COptional(BitStruct( @@ -55,7 +55,7 @@ class _Structures: "version" / Const(Int8ub, 1), "type_" / CEnum( Int8ub, - **{t.name: t.value for t in _Types} + **{t.name: t.value for t in DeviceTypes} ), "security_level" / Int8ub, "flags" / Padded(1, COptional(BitStruct( @@ -72,14 +72,13 @@ class _Structures: class Device: - Types = _Types Structures = _Structures supported_structure = Structures.v2 def __init__( self, *_: Any, - type_: Types, + type_: DeviceTypes, security_level: int, flags: Optional[dict], private_key: Optional[bytes], @@ -103,7 +102,7 @@ class Device: if not private_key: raise ValueError("Private Key is required, the WVD does not contain one or is malformed.") - self.type = self.Types[type_] if isinstance(type_, str) else type_ + self.type = DeviceTypes[type_] if isinstance(type_, str) else type_ self.security_level = security_level self.flags = flags or {} self.private_key = RSA.importKey(private_key) @@ -238,4 +237,4 @@ class Device: raise ValueError(f"Device Data seems to be corrupt or invalid, or migration failed, {e}") -__ALL__ = (Device,) +__ALL__ = (Device, DeviceTypes) diff --git a/pywidevine/main.py b/pywidevine/main.py index 702dde2..a4a2c38 100644 --- a/pywidevine/main.py +++ b/pywidevine/main.py @@ -13,7 +13,7 @@ from unidecode import UnidecodeError, unidecode from pywidevine import __version__ from pywidevine.cdm import Cdm -from pywidevine.device import Device +from pywidevine.device import Device, DeviceTypes from pywidevine.license_protocol_pb2 import FileHashes, LicenseType from pywidevine.pssh import PSSH @@ -164,7 +164,7 @@ def test(ctx: click.Context, device: Path, privacy: bool) -> None: @main.command() -@click.option("-t", "--type", "type_", type=click.Choice([x.name for x in Device.Types], case_sensitive=False), +@click.option("-t", "--type", "type_", type=click.Choice([x.name for x in DeviceTypes], case_sensitive=False), required=True, help="Device Type") @click.option("-l", "--level", type=click.IntRange(1, 3), required=True, help="Device Security Level") @click.option("-k", "--key", type=Path, required=True, help="Device RSA Private Key in PEM or DER format") @@ -195,7 +195,7 @@ def create_device( log = logging.getLogger("create-device") device = Device( - type_=Device.Types[type_.upper()], + type_=DeviceTypes[type_.upper()], security_level=level, flags=None, private_key=key.read_bytes(), diff --git a/pywidevine/remotecdm.py b/pywidevine/remotecdm.py index c3d9395..78e304d 100644 --- a/pywidevine/remotecdm.py +++ b/pywidevine/remotecdm.py @@ -12,7 +12,7 @@ from Crypto.Signature import pss from google.protobuf.message import DecodeError from pywidevine.cdm import Cdm -from pywidevine.device import Device +from pywidevine.device import Device, DeviceTypes from pywidevine.exceptions import (DeviceMismatch, InvalidInitData, InvalidLicenseMessage, InvalidLicenseType, SignatureMismatch) from pywidevine.key import Key @@ -26,7 +26,7 @@ class RemoteCdm(Cdm): def __init__( self, - device_type: Union[Device.Types, str], + device_type: Union[DeviceTypes, str], system_id: int, security_level: int, host: str, @@ -37,9 +37,9 @@ class RemoteCdm(Cdm): if not device_type: raise ValueError("Device Type must be provided") if isinstance(device_type, str): - device_type = Device.Types[device_type] - if not isinstance(device_type, Device.Types): - raise TypeError(f"Expected device_type to be a {Device.Types!r} not {device_type!r}") + device_type = DeviceTypes[device_type] + if not isinstance(device_type, DeviceTypes): + raise TypeError(f"Expected device_type to be a {DeviceTypes!r} not {device_type!r}") if not system_id: raise ValueError("System ID must be provided")