Devine
/
pywidevine
mirror of https://github.com/devine-dl/pywidevine.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

PSSH: Rework overwrite_key_ids as set_key_ids method

This commit is contained in:
rlaphoenix 2022-08-06 13:42:00 +01:00
parent fc77f064ca
commit 1064c7953c
1 changed files with 25 additions and 18 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

43
pywidevine/pssh.py
View File

@ -96,7 +96,7 @@ class PSSH:
self.version = box.version self.version = box.version
self.flags = box.flags self.flags = box.flags
self.system_id = box.system_ID self.system_id = box.system_ID
self.key_ids = box.key_IDs self.__key_ids = box.key_IDs
self.init_data = box.init_data self.init_data = box.init_data
@classmethod @classmethod
@ -175,10 +175,12 @@ class PSSH:
init_data=[init_data, b""][init_data is None] init_data=[init_data, b""][init_data is None]
))) )))
if key_ids and version == 0: pssh = cls(box)
PSSH.overwrite_key_ids(box, [UUID(bytes=x) for x in key_ids])
return cls(box) if key_ids and version == 0:
pssh.set_key_ids([UUID(bytes=x) for x in key_ids])
return pssh
@classmethod @classmethod
def from_playready_pssh(cls, box: Container) -> PSSH: def from_playready_pssh(cls, box: Container) -> PSSH:
@ -272,25 +274,30 @@ class PSSH:
raise ValueError(f"Unsupported Box {box!r}") raise ValueError(f"Unsupported Box {box!r}")
@staticmethod def set_key_ids(self, key_ids: list[UUID]) -> None:
def overwrite_key_ids(box: Container, key_ids: list[UUID]) -> Container: """Overwrite all Key IDs with the specified Key IDs."""
"""Overwrite all Key IDs in PSSH box with the specified Key IDs.""" if self.system_id != PSSH.SystemId.Widevine:
if box.system_ID != PSSH.SystemId.Widevine: # TODO: Add support for setting the Key IDs in a PlayReady Header
raise ValueError(f"Only Widevine PSSH Boxes are supported, not {box.system_ID}.") raise ValueError(f"Only Widevine PSSH Boxes are supported, not {self.system_id}.")
if box.version == 1 or box.key_IDs: if not isinstance(key_ids, list):
# only use key_IDs if version is 1, or it's already being used raise TypeError(f"Expecting key_ids to be a list, not {key_ids!r}")
if not all(isinstance(x, UUID) for x in key_ids):
not_uuid = [x for x in key_ids if not isinstance(x, UUID)]
raise TypeError(f"All Key IDs in key_ids must be a {UUID}, not {not_uuid}")
if self.version == 1 or self.__key_ids:
# only use v1 box key_ids if version is 1, or it's already being used
# this is in case the service stupidly expects it for version 0 # this is in case the service stupidly expects it for version 0
box.key_IDs = key_ids self.__key_ids = key_ids
init = WidevinePsshData() cenc_header = WidevinePsshData()
init.ParseFromString(box.init_data) cenc_header.ParseFromString(self.init_data)
init.key_ids[:] = [ cenc_header.key_ids[:] = [
key_id.bytes key_id.bytes
for key_id in key_ids for key_id in key_ids
] ]
box.init_data = init.SerializeToString() self.init_data = cenc_header.SerializeToString()
return box