Rework Cdm as a Session Key/Store Cdm

There's a few benefits to this but the main one being storage for each "request". We can now change Service Certificate per-session for example rather than for the entire Cdm object. In a multi-threaded scenario this can be a necessity more than anything.

The device is the only bit of data left that does not get stored in a session. This is mostly due to myself not seeing it being switched out often and setting it per-session would likely be cumbersome.

Some other small improvements are all around. There's a ton of doc-string improvements, typing improvements, verification of types, and there's now custom Exceptions.

In terms of bug fixes there isn't any I fixed explicitly but a possible issue in decrypt() relating the Key Labels may now be fixed.

I've moved the Keys from the return of parse_license() to the session data, with decrypt() now loading them from the session data instead. This keeps the decryption keys out of the view of the caller but it is by no way impossible to get those keys. It is incredibly trivial to access the session and get the keys from the Cdm manually.

A session limit of 50 is still set by the Cdm.
This commit is contained in:
rlaphoenix 2022-07-30 04:44:52 +01:00
parent 58186de464
commit 3536caf5f9
1 changed files with 205 additions and 89 deletions

View File

@ -1,10 +1,11 @@
import base64 import base64
import binascii
import random import random
import subprocess import subprocess
import sys import sys
import time import time
from pathlib import Path from pathlib import Path
from typing import Union, Optional from typing import Union, Container, Optional
from uuid import UUID from uuid import UUID
from Crypto.Cipher import AES, PKCS1_OAEP from Crypto.Cipher import AES, PKCS1_OAEP
@ -13,15 +14,17 @@ from Crypto.PublicKey import RSA
from Crypto.Random import get_random_bytes from Crypto.Random import get_random_bytes
from Crypto.Signature import pss from Crypto.Signature import pss
from Crypto.Util import Padding from Crypto.Util import Padding
from construct import Container
from google.protobuf.message import DecodeError from google.protobuf.message import DecodeError
from pywidevine.utils import get_binary_path
from pywidevine.license_protocol_pb2 import LicenseType, SignedMessage, LicenseRequest, ProtocolVersion, \
SignedDrmCertificate, DrmCertificate, EncryptedClientIdentification, ClientIdentification, License
from pywidevine.device import Device from pywidevine.device import Device
from pywidevine.exceptions import TooManySessions, InvalidSession, InvalidLicenseType, SignatureMismatch, \
InvalidInitData, InvalidLicenseMessage, NoKeysLoaded, InvalidContext
from pywidevine.key import Key from pywidevine.key import Key
from pywidevine.license_protocol_pb2 import DrmCertificate, SignedMessage, SignedDrmCertificate, LicenseType, \
LicenseRequest, ProtocolVersion, ClientIdentification, EncryptedClientIdentification, License
from pywidevine.pssh import PSSH from pywidevine.pssh import PSSH
from pywidevine.session import Session
from pywidevine.utils import get_binary_path
class Cdm: class Cdm:
@ -57,45 +60,47 @@ class Cdm:
root_cert = DrmCertificate() root_cert = DrmCertificate()
root_cert.ParseFromString(root_signed_cert.drm_certificate) root_cert.ParseFromString(root_signed_cert.drm_certificate)
NUM_OF_SESSIONS = 0
MAX_NUM_OF_SESSIONS = 50 # most common limit MAX_NUM_OF_SESSIONS = 50 # most common limit
def __init__(self, device: Device, init_data: Union[Container, bytes, str]): def __init__(self, device: Device):
"""Initialize a Widevine Content Decryption Module (CDM)."""
if not device:
raise ValueError("A Widevine Device must be provided.")
self.device = device
self._sessions: dict[bytes, Session] = {}
def open(self) -> bytes:
""" """
Open a Widevine Content Decryption Module (CDM) session. Open a Widevine Content Decryption Module (CDM) session.
Parameters: Raises:
device: Widevine Device containing the Client ID, Device Private Key, and TooManySessions: If the session cannot be opened as limit has been reached.
more device-specific information.
init_data: Widevine Cenc Header (Init Data) or a Protection System Specific
Header Box to take the init data from.
Devices have a limit on how many sessions can be open and active concurrently.
The limit is different for each device and security level, most commonly 50.
This limit is handled by the OEM Crypto API. Multiple sessions can be open at
a time and sessions should be closed when no longer needed.
""" """
if not device: if len(self._sessions) > self.MAX_NUM_OF_SESSIONS:
raise ValueError("A Widevine Device must be provided.") raise TooManySessions(f"Too many Sessions open ({self.MAX_NUM_OF_SESSIONS}).")
if not init_data:
raise ValueError("Init Data (or a PSSH) must be provided.")
if self.NUM_OF_SESSIONS >= self.MAX_NUM_OF_SESSIONS: session = Session()
raise ValueError( self._sessions[session.id] = session
f"Too many Sessions open {self.NUM_OF_SESSIONS}/{self.MAX_NUM_OF_SESSIONS}. "
f"Close some Sessions to be able to open more."
)
self.NUM_OF_SESSIONS += 1 return session.id
self.device = device def close(self, session_id: bytes) -> None:
self.init_data = PSSH.get_as_box(init_data).init_data """
Close a Widevine Content Decryption Module (CDM) session.
self.session_id = get_random_bytes(16) Parameters:
self.service_certificate: Optional[DrmCertificate] = None session_id: Session identifier.
self.context: dict[bytes, tuple[bytes, bytes]] = {}
def set_service_certificate(self, certificate: Union[bytes, str]) -> str: Raises:
InvalidSession: If the Session identifier is invalid.
"""
session = self._sessions.get(session_id)
if not session:
raise InvalidSession(f"Session identifier {session_id!r} is invalid.")
del self._sessions[session_id]
def set_service_certificate(self, session_id: bytes, certificate: Union[bytes, str]) -> str:
""" """
Set a Service Privacy Certificate for Privacy Mode. (optional but recommended) Set a Service Privacy Certificate for Privacy Mode. (optional but recommended)
@ -108,19 +113,31 @@ class Cdm:
containing a SignedDrmCertificate. containing a SignedDrmCertificate.
Parameters: Parameters:
session_id: Session identifier.
certificate: SignedDrmCertificate (or SignedMessage containing one) in Base64 certificate: SignedDrmCertificate (or SignedMessage containing one) in Base64
or Bytes form obtained from the Service. Some services have their own, or Bytes form obtained from the Service. Some services have their own,
but most use the common privacy cert, (common_privacy_cert). but most use the common privacy cert, (common_privacy_cert).
Raises: Raises:
InvalidSession: If the Session identifier is invalid.
DecodeError: If the certificate could not be parsed as a SignedDrmCertificate DecodeError: If the certificate could not be parsed as a SignedDrmCertificate
nor a SignedMessage containing a SignedDrmCertificate. nor a SignedMessage containing a SignedDrmCertificate.
ValueError: If the SignedDrmCertificate signature is invalid. SignatureMismatch: If the Signature of the SignedDrmCertificate does not
match the underlying DrmCertificate.
Returns the Service Provider ID of the verified DrmCertificate if successful. Returns the Service Provider ID of the verified DrmCertificate if successful.
""" """
session = self._sessions.get(session_id)
if not session:
raise InvalidSession(f"Session identifier {session_id!r} is invalid.")
if isinstance(certificate, str): if isinstance(certificate, str):
try:
certificate = base64.b64decode(certificate) # assuming base64 certificate = base64.b64decode(certificate) # assuming base64
except binascii.Error:
raise DecodeError("Could not decode certificate string as Base64, expected bytes.")
elif not isinstance(certificate, bytes):
raise DecodeError(f"Expecting Certificate to be bytes, not {certificate!r}")
signed_message = SignedMessage() signed_message = SignedMessage()
signed_drm_certificate = SignedDrmCertificate() signed_drm_certificate = SignedDrmCertificate()
@ -145,30 +162,64 @@ class Cdm:
signature=signed_drm_certificate.signature signature=signed_drm_certificate.signature
) )
except (ValueError, TypeError): except (ValueError, TypeError):
raise ValueError("Signature Mismatch on SignedDrmCertificate, rejecting certificate") raise SignatureMismatch("Signature Mismatch on SignedDrmCertificate, rejecting certificate")
else: else:
drm_certificate = DrmCertificate() drm_certificate = DrmCertificate()
drm_certificate.ParseFromString(signed_drm_certificate.drm_certificate) drm_certificate.ParseFromString(signed_drm_certificate.drm_certificate)
self.service_certificate = drm_certificate session.service_certificate = drm_certificate
return self.service_certificate.provider_id return drm_certificate.provider_id
def get_license_challenge(self, type_: Union[int, str] = LicenseType.STREAMING, privacy_mode: bool = True) -> bytes: def get_license_challenge(
self,
session_id: bytes,
init_data: Union[Container, bytes, str],
type_: Union[int, str] = LicenseType.STREAMING,
privacy_mode: bool = True
) -> bytes:
""" """
Get a License Challenge to send to a License Server. Get a License Request (Challenge) to send to a License Server.
Parameters: Parameters:
type_: Type of License you wish to exchange, often `STREAMING`. session_id: Session identifier.
The `OFFLINE` Licenses are for Offline licensing of Downloaded content. init_data: Widevine Cenc Header (Init Data) or a Protection System Specific
Header Box to take the init data from.
type_: Type of License you wish to exchange, often `STREAMING`. The `OFFLINE`
Licenses are for Offline licensing of Downloaded content.
privacy_mode: Encrypt the Client ID using the Privacy Certificate. If the privacy_mode: Encrypt the Client ID using the Privacy Certificate. If the
privacy certificate is not set yet, this does nothing. privacy certificate is not set yet, this does nothing.
Raises:
InvalidSession: If the Session identifier is invalid.
InvalidInitData: If the Init Data (or PSSH box) provided is invalid.
InvalidLicenseType: If the type_ parameter value is not a License Type. It
must be a LicenseType enum, or a string/int representing the enum's keys
or values.
Returns a SignedMessage containing a LicenseRequest message. It's signed with Returns a SignedMessage containing a LicenseRequest message. It's signed with
the Private Key of the device provision. the Private Key of the device provision.
""" """
request_id = get_random_bytes(16) session = self._sessions.get(session_id)
if not session:
raise InvalidSession(f"Session identifier {session_id!r} is invalid.")
if isinstance(type_, str): if not init_data:
raise InvalidInitData("The init_data must not be empty.")
try:
init_data = PSSH.get_as_box(init_data).init_data
except (ValueError, binascii.Error, DecodeError) as e:
raise InvalidInitData(str(e))
try:
if isinstance(type_, int):
LicenseType.Name(int(type_))
elif isinstance(type_, str):
type_ = LicenseType.Value(type_) type_ = LicenseType.Value(type_)
elif not isinstance(type_, LicenseType):
raise InvalidLicenseType()
except ValueError:
raise InvalidLicenseType(f"License Type {type_!r} is invalid")
request_id = get_random_bytes(16)
license_request = LicenseRequest() license_request = LicenseRequest()
license_request.type = LicenseRequest.RequestType.Value("NEW") license_request.type = LicenseRequest.RequestType.Value("NEW")
@ -176,49 +227,76 @@ class Cdm:
license_request.protocol_version = ProtocolVersion.Value("VERSION_2_1") license_request.protocol_version = ProtocolVersion.Value("VERSION_2_1")
license_request.key_control_nonce = random.randrange(1, 2 ** 31) license_request.key_control_nonce = random.randrange(1, 2 ** 31)
license_request.content_id.widevine_pssh_data.pssh_data.append(self.init_data) license_request.content_id.widevine_pssh_data.pssh_data.append(init_data)
license_request.content_id.widevine_pssh_data.license_type = type_ license_request.content_id.widevine_pssh_data.license_type = type_
license_request.content_id.widevine_pssh_data.request_id = request_id license_request.content_id.widevine_pssh_data.request_id = request_id
if self.service_certificate and privacy_mode: if session.service_certificate and privacy_mode:
# encrypt the client id for privacy mode # encrypt the client id for privacy mode
license_request.encrypted_client_id.CopyFrom(self.encrypt_client_id( license_request.encrypted_client_id.CopyFrom(self.encrypt_client_id(
client_id=self.device.client_id, client_id=self.device.client_id,
service_certificate=self.service_certificate service_certificate=session.service_certificate
)) ))
else: else:
license_request.client_id.CopyFrom(self.device.client_id) license_request.client_id.CopyFrom(self.device.client_id)
license_message = SignedMessage() license_message = SignedMessage()
license_message.type = SignedMessage.MessageType.Value("LICENSE_REQUEST") license_message.type = SignedMessage.MessageType.LICENSE_REQUEST
license_message.msg = license_request.SerializeToString() license_message.msg = license_request.SerializeToString()
license_message.signature = pss. \ license_message.signature = pss. \
new(self.device.private_key). \ new(self.device.private_key). \
sign(SHA1.new(license_message.msg)) sign(SHA1.new(license_message.msg))
self.context[request_id] = self.derive_context(license_message.msg) session.context[request_id] = self.derive_context(license_message.msg)
return license_message.SerializeToString() return license_message.SerializeToString()
def parse_license(self, license_message: Union[bytes, str]) -> list[Key]: def parse_license(self, session_id: bytes, license_message: Union[SignedMessage, bytes, str]) -> None:
"""
Load Keys from a License Message from a License Server Response.
Parameters:
session_id: Session identifier.
license_message: A SignedMessage containing a License message.
Raises:
InvalidSession: If the Session identifier is invalid.
InvalidLicenseMessage: The License message could not be decoded as a Signed
Message or License message.
InvalidContext: If the Session has no Context Data. This is likely to happen
if the License Challenge was not made by this CDM instance, or was not
by this CDM at all. It could also happen if the Session is closed after
calling parse_license but not before it got the context data.
SignatureMismatch: If the Signature of the License SignedMessage does not
match the underlying License.
"""
session = self._sessions.get(session_id)
if not session:
raise InvalidSession(f"Session identifier {session_id!r} is invalid.")
if not license_message: if not license_message:
raise ValueError("Cannot parse an empty license_message as a SignedMessage") raise InvalidLicenseMessage("Cannot parse an empty license_message")
if isinstance(license_message, str): if isinstance(license_message, str):
try:
license_message = base64.b64decode(license_message) license_message = base64.b64decode(license_message)
except (binascii.Error, binascii.Incomplete) as e:
raise InvalidLicenseMessage(f"Could not decode license_message as Base64, {e}")
if isinstance(license_message, bytes): if isinstance(license_message, bytes):
signed_message = SignedMessage() signed_message = SignedMessage()
try: try:
signed_message.ParseFromString(license_message) signed_message.ParseFromString(license_message)
except DecodeError: except DecodeError as e:
raise ValueError("Failed to parse license_message as a SignedMessage") raise InvalidLicenseMessage(f"Could not parse license_message as a SignedMessage, {e}")
license_message = signed_message license_message = signed_message
if not isinstance(license_message, SignedMessage): if not isinstance(license_message, SignedMessage):
raise ValueError(f"Expecting license_response to be a SignedMessage, got {license_message!r}") raise InvalidLicenseMessage(f"Expecting license_response to be a SignedMessage, got {license_message!r}")
if license_message.type != SignedMessage.MessageType.LICENSE: if license_message.type != SignedMessage.MessageType.LICENSE:
raise ValueError( raise InvalidLicenseMessage(
f"Expecting a LICENSE message, not a " f"Expecting a LICENSE message, not a "
f"'{SignedMessage.MessageType.Name(license_message.type)}' message." f"'{SignedMessage.MessageType.Name(license_message.type)}' message."
) )
@ -226,9 +304,9 @@ class Cdm:
licence = License() licence = License()
licence.ParseFromString(license_message.msg) licence.ParseFromString(license_message.msg)
context = self.context.get(licence.id.request_id) context = session.context.get(licence.id.request_id)
if not context: if not context:
raise ValueError("Cannot parse a license message without first making a license request") raise InvalidContext("Cannot parse a license message without first making a license request")
session_key = PKCS1_OAEP. \ session_key = PKCS1_OAEP. \
new(self.device.private_key). \ new(self.device.private_key). \
@ -242,60 +320,97 @@ class Cdm:
digest() digest()
if license_message.signature != computed_signature: if license_message.signature != computed_signature:
raise ValueError("Signature Mismatch on License Message, rejecting license") raise SignatureMismatch("Signature Mismatch on License Message, rejecting license")
return [ session.keys = [
Key.from_key_container(key, enc_key) Key.from_key_container(key, enc_key)
for key in licence.key for key in licence.key
] ]
@staticmethod def decrypt(
def decrypt(content_keys: dict[UUID, str], input_: Path, output: Path, temp: Optional[Path] = None): self,
session_id: bytes,
input_file: Union[Path, str],
output_file: Union[Path, str],
temp_dir: Optional[Union[Path, str]] = None,
exists_ok: bool = False
):
""" """
Decrypt a Widevine-encrypted file using Shaka-packager. Decrypt a Widevine-encrypted file using Shaka-packager.
Shaka-packager is much more stable than mp4decrypt. Shaka-packager is much more stable than mp4decrypt.
Parameters:
session_id: Session identifier.
input_file: File to be decrypted with Session's currently loaded keys.
output_file: Location to save decrypted file.
temp_dir: Directory to store temporary data while decrypting.
exists_ok: Allow overwriting the output_file if it exists.
Raises: Raises:
EnvironmentError if the Shaka Packager executable could not be found. ValueError: If the input or output paths have not been supplied or are
ValueError if the track has not yet been downloaded. invalid.
SubprocessError if Shaka Packager returned a non-zero exit code. FileNotFoundError: If the input file path does not exist.
FileExistsError: If the output file path already exists. Ignored if exists_ok
is set to True.
NoKeysLoaded: No License was parsed for this Session, No Keys available.
EnvironmentError: If the shaka-packager executable could not be found.
subprocess.CalledProcessError: If the shaka-packager call returned a non-zero
exit code.
""" """
if not content_keys: if not input_file:
raise ValueError("Cannot decrypt without any Content Keys")
if not input_:
raise ValueError("Cannot decrypt nothing, specify an input path") raise ValueError("Cannot decrypt nothing, specify an input path")
if not output: if not output_file:
raise ValueError("Cannot decrypt nowhere, specify an output path") raise ValueError("Cannot decrypt nowhere, specify an output path")
if not isinstance(input_file, (Path, str)):
raise ValueError(f"Expecting input_file to be a Path or str, got {input_file!r}")
if not isinstance(output_file, (Path, str)):
raise ValueError(f"Expecting output_file to be a Path or str, got {output_file!r}")
if not isinstance(temp_dir, (Path, str)) and temp_dir is not None:
raise ValueError(f"Expecting temp_dir to be a Path or str, got {temp_dir!r}")
input_file = Path(input_file)
output_file = Path(output_file)
if temp_dir:
temp_dir = Path(temp_dir)
if not input_file.is_file():
raise FileNotFoundError(f"Input file does not exist, {input_file}")
if output_file.is_file() and not exists_ok:
raise FileExistsError(f"Output file already exists, {output_file}")
session = self._sessions.get(session_id)
if not session:
raise InvalidSession(f"Session identifier {session_id!r} is invalid.")
if not session.keys:
raise NoKeysLoaded("No Keys are loaded yet, cannot decrypt")
platform = {"win32": "win", "darwin": "osx"}.get(sys.platform, sys.platform) platform = {"win32": "win", "darwin": "osx"}.get(sys.platform, sys.platform)
executable = get_binary_path("shaka-packager", f"packager-{platform}", f"packager-{platform}-x64") executable = get_binary_path("shaka-packager", f"packager-{platform}", f"packager-{platform}-x64")
if not executable: if not executable:
raise EnvironmentError("Shaka Packager executable not found but is required") raise EnvironmentError("Shaka Packager executable not found but is required")
args = [ args = [
f"input={input_},stream=0,output={output}", f"input={input_file},stream=0,output={output_file}",
"--enable_raw_key_decryption", "--keys", "--enable_raw_key_decryption",
",".join([ "--keys", ",".join([
*[ label
f"label={i}:key_id={kid.hex}:key={key.lower()}" for i, key in enumerate(session.keys)
for i, (kid, key) in enumerate(content_keys.items()) for label in [
], f"label=1_{i}:key_id={key.kid.hex}:key={key.key.hex()}",
*[ # some services need the KID blanked, e.g., Apple TV+
# Apple TV+ needs this as their files do not use the KID supplied in the manifest f"label=2_{i}:key_id={'0' * 32}:key={key.key.hex()}"
f"label={i}:key_id=00000000000000000000000000000000:key={key.lower()}"
for i, (kid, key) in enumerate(content_keys.items(), len(content_keys))
] ]
if key.type == "CONTENT"
]) ])
] ]
if temp: if temp_dir:
temp.mkdir(parents=True, exist_ok=True) temp_dir.mkdir(parents=True, exist_ok=True)
args.extend(["--temp_dir", temp]) args.extend(["--temp_dir", temp_dir])
try:
subprocess.check_call([executable, *args]) subprocess.check_call([executable, *args])
except subprocess.CalledProcessError as e:
raise subprocess.SubprocessError(f"Failed to Decrypt! Shaka Packager Error: {e}")
@staticmethod @staticmethod
def encrypt_client_id( def encrypt_client_id(
@ -365,7 +480,8 @@ class Cdm:
""" """
def _derive(session_key: bytes, context: bytes, counter: int) -> bytes: def _derive(session_key: bytes, context: bytes, counter: int) -> bytes:
return CMAC.new(session_key, ciphermod=AES). \ return CMAC. \
new(session_key, ciphermod=AES). \
update(counter.to_bytes(1, "big") + context). \ update(counter.to_bytes(1, "big") + context). \
digest() digest()