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.
This commit is contained in:
rlaphoenix 2023-11-08 22:38:32 +00:00
parent c362192c11
commit 79cdbc007c
4 changed files with 20 additions and 21 deletions

View File

@ -18,7 +18,7 @@ from Crypto.Signature import pss
from Crypto.Util import Padding from Crypto.Util import Padding
from google.protobuf.message import DecodeError 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, from pywidevine.exceptions import (InvalidContext, InvalidInitData, InvalidLicenseMessage, InvalidLicenseType,
InvalidSession, NoKeysLoaded, SignatureMismatch, TooManySessions) InvalidSession, NoKeysLoaded, SignatureMismatch, TooManySessions)
from pywidevine.key import Key from pywidevine.key import Key
@ -66,7 +66,7 @@ class Cdm:
def __init__( def __init__(
self, self,
device_type: Union[Device.Types, str], device_type: Union[DeviceTypes, str],
system_id: int, system_id: int,
security_level: int, security_level: int,
client_id: ClientIdentification, client_id: ClientIdentification,
@ -76,9 +76,9 @@ class Cdm:
if not device_type: if not device_type:
raise ValueError("Device Type must be provided") raise ValueError("Device Type must be provided")
if isinstance(device_type, str): if isinstance(device_type, str):
device_type = Device.Types[device_type] device_type = DeviceTypes[device_type]
if not isinstance(device_type, Device.Types): if not isinstance(device_type, DeviceTypes):
raise TypeError(f"Expected device_type to be a {Device.Types!r} not {device_type!r}") raise TypeError(f"Expected device_type to be a {DeviceTypes!r} not {device_type!r}")
if not system_id: if not system_id:
raise ValueError("System ID must be provided") raise ValueError("System ID must be provided")
@ -306,7 +306,7 @@ class Cdm:
f"Available values: {LicenseType.keys()}" 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 # 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 # Bytes 5-8 does not seem random, in real tests they have been consecutive \x00 or \xFF
# Real example: A0DCE548000000000500000000000000 # Real example: A0DCE548000000000500000000000000

View File

@ -17,7 +17,7 @@ from google.protobuf.message import DecodeError
from pywidevine.license_protocol_pb2 import ClientIdentification, DrmCertificate, FileHashes, SignedDrmCertificate from pywidevine.license_protocol_pb2 import ClientIdentification, DrmCertificate, FileHashes, SignedDrmCertificate
class _Types(Enum): class DeviceTypes(Enum):
CHROME = 1 CHROME = 1
ANDROID = 2 ANDROID = 2
@ -36,7 +36,7 @@ class _Structures:
"version" / Const(Int8ub, 2), "version" / Const(Int8ub, 2),
"type_" / CEnum( "type_" / CEnum(
Int8ub, Int8ub,
**{t.name: t.value for t in _Types} **{t.name: t.value for t in DeviceTypes}
), ),
"security_level" / Int8ub, "security_level" / Int8ub,
"flags" / Padded(1, COptional(BitStruct( "flags" / Padded(1, COptional(BitStruct(
@ -55,7 +55,7 @@ class _Structures:
"version" / Const(Int8ub, 1), "version" / Const(Int8ub, 1),
"type_" / CEnum( "type_" / CEnum(
Int8ub, Int8ub,
**{t.name: t.value for t in _Types} **{t.name: t.value for t in DeviceTypes}
), ),
"security_level" / Int8ub, "security_level" / Int8ub,
"flags" / Padded(1, COptional(BitStruct( "flags" / Padded(1, COptional(BitStruct(
@ -72,14 +72,13 @@ class _Structures:
class Device: class Device:
Types = _Types
Structures = _Structures Structures = _Structures
supported_structure = Structures.v2 supported_structure = Structures.v2
def __init__( def __init__(
self, self,
*_: Any, *_: Any,
type_: Types, type_: DeviceTypes,
security_level: int, security_level: int,
flags: Optional[dict], flags: Optional[dict],
private_key: Optional[bytes], private_key: Optional[bytes],
@ -103,7 +102,7 @@ class Device:
if not private_key: if not private_key:
raise ValueError("Private Key is required, the WVD does not contain one or is malformed.") 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.security_level = security_level
self.flags = flags or {} self.flags = flags or {}
self.private_key = RSA.importKey(private_key) 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}") raise ValueError(f"Device Data seems to be corrupt or invalid, or migration failed, {e}")
__ALL__ = (Device,) __ALL__ = (Device, DeviceTypes)

View File

@ -13,7 +13,7 @@ from unidecode import UnidecodeError, unidecode
from pywidevine import __version__ from pywidevine import __version__
from pywidevine.cdm import Cdm 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.license_protocol_pb2 import FileHashes, LicenseType
from pywidevine.pssh import PSSH from pywidevine.pssh import PSSH
@ -164,7 +164,7 @@ def test(ctx: click.Context, device: Path, privacy: bool) -> None:
@main.command() @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") required=True, help="Device Type")
@click.option("-l", "--level", type=click.IntRange(1, 3), required=True, help="Device Security Level") @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") @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") log = logging.getLogger("create-device")
device = Device( device = Device(
type_=Device.Types[type_.upper()], type_=DeviceTypes[type_.upper()],
security_level=level, security_level=level,
flags=None, flags=None,
private_key=key.read_bytes(), private_key=key.read_bytes(),

View File

@ -12,7 +12,7 @@ from Crypto.Signature import pss
from google.protobuf.message import DecodeError from google.protobuf.message import DecodeError
from pywidevine.cdm import Cdm from pywidevine.cdm import Cdm
from pywidevine.device import Device from pywidevine.device import Device, DeviceTypes
from pywidevine.exceptions import (DeviceMismatch, InvalidInitData, InvalidLicenseMessage, InvalidLicenseType, from pywidevine.exceptions import (DeviceMismatch, InvalidInitData, InvalidLicenseMessage, InvalidLicenseType,
SignatureMismatch) SignatureMismatch)
from pywidevine.key import Key from pywidevine.key import Key
@ -26,7 +26,7 @@ class RemoteCdm(Cdm):
def __init__( def __init__(
self, self,
device_type: Union[Device.Types, str], device_type: Union[DeviceTypes, str],
system_id: int, system_id: int,
security_level: int, security_level: int,
host: str, host: str,
@ -37,9 +37,9 @@ class RemoteCdm(Cdm):
if not device_type: if not device_type:
raise ValueError("Device Type must be provided") raise ValueError("Device Type must be provided")
if isinstance(device_type, str): if isinstance(device_type, str):
device_type = Device.Types[device_type] device_type = DeviceTypes[device_type]
if not isinstance(device_type, Device.Types): if not isinstance(device_type, DeviceTypes):
raise TypeError(f"Expected device_type to be a {Device.Types!r} not {device_type!r}") raise TypeError(f"Expected device_type to be a {DeviceTypes!r} not {device_type!r}")
if not system_id: if not system_id:
raise ValueError("System ID must be provided") raise ValueError("System ID must be provided")