From abf6c71688b79a68c08be560b825acfd043e3ac5 Mon Sep 17 00:00:00 2001
From: rlaphoenix <rlaphoenix@pm.me>
Date: Wed, 8 Mar 2023 22:45:41 +0000
Subject: [PATCH] Specify HLS Track Key IDs to prepare_drm

This also moves the init data code before drm related code, just so it has the init data ready to retrieve the Key ID from.
---
 devine/core/manifests/hls.py | 67 ++++++++++++++++++------------------
 1 file changed, 34 insertions(+), 33 deletions(-)

diff --git a/devine/core/manifests/hls.py b/devine/core/manifests/hls.py
index 4462790..c99007a 100644
--- a/devine/core/manifests/hls.py
+++ b/devine/core/manifests/hls.py
@@ -221,42 +221,11 @@ class HLS:
                 # the track already started downloading, but another failed or was stopped
                 raise KeyboardInterrupt()
 
-            segment_save_path = (save_dir / filename).with_suffix(".mp4")
-
-            with drm_lock:
-                newest_segment_key = segment_key.get()
-                try:
-                    if segment.key and newest_segment_key[1] != segment.key:
-                        try:
-                            drm = HLS.get_drm(
-                                # TODO: We append master.keys because m3u8 class only puts the last EXT-X-KEY
-                                #       to the segment.key property, not supporting multi-drm scenarios.
-                                #       By re-adding every single EXT-X-KEY found, we can at least try to get
-                                #       a suitable key. However, it may not match the right segment/timeframe!
-                                #       It will try to use the first key provided where possible.
-                                keys=[segment.key] + master.keys,
-                                proxy=proxy
-                            )
-                        except NotImplementedError as e:
-                            log.error(str(e))
-                            sys.exit(1)
-                        else:
-                            if drm:
-                                track.drm = drm
-                                drm = drm[0]  # just use the first supported DRM system for now
-                                log.debug("Got segment key, %s", drm)
-                                if isinstance(drm, Widevine):
-                                    # license and grab content keys
-                                    if not license_widevine:
-                                        raise ValueError("license_widevine func must be supplied to use Widevine DRM")
-                                    license_widevine(drm)
-                                newest_segment_key = (drm, segment.key)
-                finally:
-                    segment_key.put(newest_segment_key)
-
             if callable(track.OnSegmentFilter) and track.OnSegmentFilter(segment):
                 return 0
 
+            segment_save_path = (save_dir / filename).with_suffix(".mp4")
+
             newest_init_data = init_data.get()
             try:
                 if segment.init_section and (not newest_init_data or segment.discontinuity):
@@ -284,6 +253,38 @@ class HLS:
             finally:
                 init_data.put(newest_init_data)
 
+            with drm_lock:
+                newest_segment_key = segment_key.get()
+                try:
+                    if segment.key and newest_segment_key[1] != segment.key:
+                        try:
+                            drm = HLS.get_drm(
+                                # TODO: We append master.keys because m3u8 class only puts the last EXT-X-KEY
+                                #       to the segment.key property, not supporting multi-drm scenarios.
+                                #       By re-adding every single EXT-X-KEY found, we can at least try to get
+                                #       a suitable key. However, it may not match the right segment/timeframe!
+                                #       It will try to use the first key provided where possible.
+                                keys=[segment.key] + master.keys,
+                                proxy=proxy
+                            )
+                        except NotImplementedError as e:
+                            log.error(str(e))
+                            sys.exit(1)
+                        else:
+                            if drm:
+                                track.drm = drm
+                                drm = drm[0]  # just use the first supported DRM system for now
+                                log.debug("Got segment key, %s", drm)
+                                if isinstance(drm, Widevine):
+                                    # license and grab content keys
+                                    track_kid = track.get_key_id(newest_init_data)
+                                    if not license_widevine:
+                                        raise ValueError("license_widevine func must be supplied to use Widevine DRM")
+                                    license_widevine(drm, track_kid=track_kid)
+                                newest_segment_key = (drm, segment.key)
+                finally:
+                    segment_key.put(newest_segment_key)
+
             if not segment.uri.startswith(segment.base_uri):
                 segment.uri = segment.base_uri + segment.uri