From 73bd17ec94c11bcab740a32e0f988c8726b735be Mon Sep 17 00:00:00 2001
From: rlaphoenix <rlaphoenix@pm.me>
Date: Wed, 8 Mar 2023 22:33:19 +0000
Subject: [PATCH] Implement new get_key_id() method to Track

---
 devine/core/tracks/track.py | 42 ++++++++++++++++++++++++++++++++++++-
 1 file changed, 41 insertions(+), 1 deletion(-)

diff --git a/devine/core/tracks/track.py b/devine/core/tracks/track.py
index 384928a..e0b5298 100644
--- a/devine/core/tracks/track.py
+++ b/devine/core/tracks/track.py
@@ -1,16 +1,19 @@
+import base64
 import re
 import shutil
 import subprocess
 from enum import Enum
 from pathlib import Path
 from typing import Any, Callable, Iterable, Optional, Union
+from uuid import UUID
 
 import requests
 from langcodes import Language
 
 from devine.core.constants import TERRITORY_MAP
 from devine.core.drm import DRM_T
-from devine.core.utilities import get_binary_path
+from devine.core.utilities import get_binary_path, get_boxes
+from devine.core.utils.subprocess import ffprobe
 
 
 class Track:
@@ -84,6 +87,43 @@ class Track:
             extra_parts.append(TERRITORY_MAP.get(territory, territory))
         return ", ".join(extra_parts) or None
 
+    def get_key_id(self, init_data: Optional[bytes] = None, *args, **kwargs) -> Optional[UUID]:
+        """
+        Probe the DRM encryption Key ID (KID) for this specific track.
+
+        It currently supports finding the Key ID by probing the track's stream
+        with ffprobe for `enc_key_id` data, as well as for mp4 `tenc` (Track
+        Encryption) boxes.
+
+        It explicitly ignores PSSH information like the `PSSH` box, as the box
+        is likely to contain multiple Key IDs that may or may not be for this
+        specific track.
+
+        To retrieve the initialization segment, this method calls :meth:`get_init_segment`
+        with the positional and keyword arguments. The return value of `get_init_segment`
+        is then used to determine the Key ID.
+
+        Returns:
+            The Key ID as a UUID object, or None if the Key ID could not be determined.
+        """
+        if not init_data:
+            init_data = self.get_init_segment(*args, **kwargs)
+        if not isinstance(init_data, bytes):
+            raise TypeError(f"Expected init_data to be bytes, not {init_data!r}")
+
+        # try get via ffprobe, needed for non mp4 data e.g. WEBM from Google Play
+        probe = ffprobe(init_data)
+        if probe:
+            for stream in probe.get("streams") or []:
+                enc_key_id = stream.get("tags", {}).get("enc_key_id")
+                if enc_key_id:
+                    return UUID(bytes=base64.b64decode(enc_key_id))
+
+        # look for track encryption mp4 boxes
+        for tenc in get_boxes(init_data, b"tenc"):
+            if tenc.key_ID.int != 0:
+                return tenc.key_ID
+
     def get_init_segment(
         self,
         maximum_size: int = 20000,