PSSH: Ensure key IDs are UUIDs instead of Bytes

This reduces code duplication when actually using those key_ids.
This commit is contained in:
rlaphoenix 2022-11-18 09:00:14 +00:00
parent 2d2359f9a2
commit 7a993206a1
1 changed files with 15 additions and 14 deletions

View File

@ -134,22 +134,23 @@ class PSSH:
raise ValueError("Version 1 PSSH boxes must use either init_data or key_ids but neither were provided") raise ValueError("Version 1 PSSH boxes must use either init_data or key_ids but neither were provided")
if key_ids is not None: if key_ids is not None:
# ensure key_ids are bytes, supports hex, base64, and bytes # ensure key_ids are UUID, supports hex, base64, and bytes
key_ids = [ if not all(isinstance(x, (UUID, bytes, str)) for x in key_ids):
( not_bytes = [x for x in key_ids if not isinstance(x, (UUID, bytes, str))]
x.bytes if isinstance(x, UUID) else
bytes.fromhex(x) if all(c in string.hexdigits for c in x) else
base64.b64decode(x) if isinstance(x, str) else
x
)
for x in key_ids
]
if not all(isinstance(x, bytes) for x in key_ids):
not_bytes = [x for x in key_ids if not isinstance(x, bytes)]
raise TypeError( raise TypeError(
"Expected all of key_ids to be a UUID, hex, base64, or bytes, but one or more are not, " "Expected all of key_ids to be a UUID, hex, base64, or bytes, but one or more are not, "
f"{not_bytes!r}" f"{not_bytes!r}"
) )
key_ids = [
UUID(bytes=key_id_b)
for key_id in key_ids
for key_id_b in [
key_id.bytes if isinstance(key_id, UUID) else
bytes.fromhex(key_id) if all(c in string.hexdigits for c in key_id) else
base64.b64decode(key_id) if isinstance(key_id, str) else
key_id
]
]
if init_data is not None: if init_data is not None:
if isinstance(init_data, WidevinePsshData): if isinstance(init_data, WidevinePsshData):
@ -169,14 +170,14 @@ class PSSH:
version=version, version=version,
flags=flags, flags=flags,
system_ID=PSSH.SystemId.Widevine, system_ID=PSSH.SystemId.Widevine,
key_IDs=[UUID(bytes=kid) for kid in key_ids] if key_ids else None, key_IDs=key_ids if key_ids else None,
init_data=[init_data, b""][init_data is None] init_data=[init_data, b""][init_data is None]
))) )))
pssh = cls(box) pssh = cls(box)
if key_ids and version == 0: if key_ids and version == 0:
pssh.set_key_ids([UUID(bytes=x) for x in key_ids]) pssh.set_key_ids(key_ids)
return pssh return pssh