From 7a993206a1595afddc46767aa7512bc2361b83f3 Mon Sep 17 00:00:00 2001
From: rlaphoenix <rlaphoenix@pm.me>
Date: Fri, 18 Nov 2022 09:00:14 +0000
Subject: [PATCH] PSSH: Ensure key IDs are UUIDs instead of Bytes

This reduces code duplication when actually using those key_ids.
---
 pywidevine/pssh.py | 29 +++++++++++++++--------------
 1 file changed, 15 insertions(+), 14 deletions(-)

diff --git a/pywidevine/pssh.py b/pywidevine/pssh.py
index 9dc3a55..e7544f2 100644
--- a/pywidevine/pssh.py
+++ b/pywidevine/pssh.py
@@ -134,22 +134,23 @@ class PSSH:
                 raise ValueError("Version 1 PSSH boxes must use either init_data or key_ids but neither were provided")
 
         if key_ids is not None:
-            # ensure key_ids are bytes, supports hex, base64, and bytes
-            key_ids = [
-                (
-                    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)]
+            # ensure key_ids are UUID, supports hex, base64, and bytes
+            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))]
                 raise TypeError(
                     "Expected all of key_ids to be a UUID, hex, base64, or bytes, but one or more are not, "
                     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 isinstance(init_data, WidevinePsshData):
@@ -169,14 +170,14 @@ class PSSH:
             version=version,
             flags=flags,
             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]
         )))
 
         pssh = cls(box)
 
         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