From 1b40c2b36945cccac606f3304e70d12cc6d374f4 Mon Sep 17 00:00:00 2001
From: rlaphoenix <rlaphoenix@pm.me>
Date: Fri, 18 Nov 2022 09:40:55 +0000
Subject: [PATCH] PSSH: Set Key IDs more effectively via set_key_ids()

This reduces reading complexity of why and when pssh.set_key_ids() was being run. Generally less code repetition effectively.
---
 pywidevine/pssh.py | 15 +++++++++------
 1 file changed, 9 insertions(+), 6 deletions(-)

diff --git a/pywidevine/pssh.py b/pywidevine/pssh.py
index 4c50368..e9167fc 100644
--- a/pywidevine/pssh.py
+++ b/pywidevine/pssh.py
@@ -165,18 +165,21 @@ class PSSH:
                     f"Expecting init_data to be {WidevinePsshData}, hex, base64, or bytes, not {init_data!r}"
                 )
 
-        box = Box.parse(Box.build(dict(
+        pssh = cls(Box.parse(Box.build(dict(
             type=b"pssh",
             version=version,
             flags=flags,
             system_ID=PSSH.SystemId.Widevine,
-            key_IDs=key_ids if key_ids and version == 1 else None,
             init_data=[init_data, b""][init_data is None]
-        )))
+            # key_IDs should not be set yet
+        ))))
 
-        pssh = cls(box)
-
-        if key_ids and version == 0:
+        if key_ids:
+            # We must reinforce the version because pymp4 forces v0 if key_IDs is not set.
+            # The set_key_ids() func will set it efficiently in both init_data and the box where needed.
+            # The version must be reinforced ONLY if we have key_id data or there's a possibility of making
+            # a v1 PSSH box, that did not have key_IDs set in the PSSH box.
+            pssh.version = version
             pssh.set_key_ids(key_ids)
 
         return pssh