Map default_KID to kid:key

Author: Arrowar

StreamingCommunity/Lib/DASH/cdm_helpher.py

@@ -33,53 +33,6 @@ def filter_valid_keys(content_keys: list) -> list:
     return valid_keys
 
 
-def select_best_key(valid_keys: list) -> dict:
-    """
-    Select the best key from valid keys based on heuristics.
-    
-    Args:
-        valid_keys (list): List of valid key dictionaries
-    """
-    if len(valid_keys) == 1:
-        return valid_keys[0]
-    
-    # Heuristics for key selection:
-    # 1. Prefer keys that are not all the same character
-    # 2. Prefer keys with more entropy (variety in hex characters)
-    scored_keys = []
-    for key_info in valid_keys:
-        key_value = key_info.get('key', '')
-        score = 0
-        
-        # Score based on character variety
-        unique_chars = len(set(key_value))
-        score += unique_chars
-        
-        # Penalize keys with too many repeated patterns
-        if len(key_value) > 8:
-            
-            # Check for repeating patterns
-            has_pattern = False
-            for i in range(2, len(key_value) // 2):
-                pattern = key_value[:i]
-                if key_value.startswith(pattern * (len(key_value) // i)):
-                    has_pattern = True
-                    break
-            
-            if not has_pattern:
-                score += 10
-        
-        scored_keys.append((score, key_info))
-        console.log(f"[cyan]Found key [red]{key_info.get('kid', 'unknown')}[cyan] with score: [red]{score}")
-    
-    # Sort by score (descending) and return the best key
-    scored_keys.sort(key=lambda x: x[0], reverse=True)
-    best_key = scored_keys[0][1]
-    
-    console.log(f"[cyan]Selected key: [red]{best_key.get('kid', 'unknown')}[cyan] with score: [red]{scored_keys[0][0]}")
-    return best_key
-
-
 def get_widevine_keys(pssh: str, license_url: str, cdm_device_path: str, headers: dict = None, query_params: dict =None, key: str=None):
     """
     Extract Widevine CONTENT keys (KID/KEY) from a license using pywidevine.
@@ -174,24 +127,14 @@ def get_widevine_keys(pssh: str, license_url: str, cdm_device_path: str, headers
                 console.print("[yellow]⚠️ No CONTENT keys found in license.")
                 return None
 
-            # Filter and select the best key
+            # Filter valid keys
             valid_keys = filter_valid_keys(content_keys)
 
-            # Select the best key automatically
-            best_key = select_best_key(valid_keys)
-            
-            if best_key:
-                console.log(f"[cyan]Selected KID: [green]{best_key['kid']} [white]| [cyan]KEY: [green]{best_key['key']}")
-                
-                # Return all valid keys but with the best one first
-                result_keys = [best_key]
-                for key_info in valid_keys:
-                    if key_info['kid'] != best_key['kid']:
-                        result_keys.append(key_info)
-                
-                return result_keys
+            if valid_keys:
+                console.log(f"[cyan]Found {len(valid_keys)} valid keys for testing")
+                return valid_keys
             else:
-                console.print("[red]❌ Could not select best key")
+                console.print("[red]❌ No valid keys found")
                 return None
         else:
             content_keys = []
@@ -210,6 +153,40 @@ def get_widevine_keys(pssh: str, license_url: str, cdm_device_path: str, headers
         cdm.close(session_id)
 
 
+def map_keys_to_representations(keys: list, representations: list) -> dict:
+    """
+    Map decryption keys to representations based on their default_KID.
+    
+    Args:
+        keys (list): List of key dictionaries with 'kid' and 'key' fields
+        representations (list): List of representation dictionaries with 'default_kid' field
+    
+    Returns:
+        dict: Mapping of representation type to key info
+    """
+    key_mapping = {}
+    
+    for rep in representations:
+        rep_type = rep.get('type', 'unknown')
+        default_kid = rep.get('default_kid')
+        
+        if not default_kid:
+            continue
+            
+        for key_info in keys:
+            if key_info['kid'].lower() == default_kid.lower():
+                key_mapping[rep_type] = {
+                    'kid': key_info['kid'],
+                    'key': key_info['key'],
+                    'representation_id': rep.get('id'),
+                    'default_kid': default_kid
+                }
+                #console.log(f"[cyan]Mapped {rep_type} representation [yellow]{rep.get('id')} [cyan]to key: [red]{key_info['kid']}")
+                break
+    
+    return key_mapping
+
+
 def get_info_wvd(cdm_device_path):
     """
     Extract device information from a Widevine CDM device file (.wvd).

StreamingCommunity/Lib/DASH/downloader.py

@@ -20,7 +20,7 @@
 from .parser import MPD_Parser
 from .segments import MPD_Segments
 from .decrypt import decrypt_with_mp4decrypt
-from .cdm_helpher import get_widevine_keys
+from .cdm_helpher import get_widevine_keys, map_keys_to_representations
 
 
 # FFmpeg functions
@@ -32,7 +32,6 @@
 DOWNLOAD_SPECIFIC_SUBTITLE = config_manager.get_list('M3U8_DOWNLOAD', 'specific_list_subtitles')
 MERGE_SUBTITLE = config_manager.get_bool('M3U8_DOWNLOAD', 'merge_subs')
 CLEANUP_TMP = config_manager.get_bool('M3U8_DOWNLOAD', 'cleanup_tmp_folder')
-RETRY_LIMIT = config_manager.get_int('REQUESTS', 'max_retry')
 EXTENSION_OUTPUT = config_manager.get("M3U8_CONVERSION", "extension")
 
 
@@ -160,7 +159,7 @@ def get_representation_by_type(self, typ):
 
     def download_subtitles(self) -> bool:
         """
-        Download subtitle files based on parser's selected subtitles with retry mechanism.
+        Download subtitle files based on parser's selected subtitles.
         Returns True if successful or if no subtitles to download, False on critical error.
         """
         if not self.selected_subs or self.mpd_sub_list is None:
@@ -193,12 +192,11 @@ def download_subtitles(self) -> bool:
 
     def download_and_decrypt(self, custom_headers=None, query_params=None, key=None) -> bool:
         """
-        Download and decrypt video/audio streams. Skips download if file already exists.
+        Download and decrypt video/audio streams using automatic key mapping based on default_KID.
 
         Args:
             - custom_headers (dict): Optional HTTP headers for the license request.
             - query_params (dict): Optional query parameters to append to the license URL.
-            - license_data (str/bytes): Optional raw license data to bypass HTTP request.
             - key (str): Optional raw license data to bypass HTTP request.
         """
         if self.file_already_exists:
@@ -228,12 +226,26 @@ def download_and_decrypt(self, custom_headers=None, query_params=None, key=None)
             console.print("[red]No keys found, cannot proceed with download.")
             return False
 
+        # Map keys to representations based on default_KID
+        key_mapping = map_keys_to_representations(keys, self.parser.representations)
+        
+        if not key_mapping:
+            console.print("[red]Could not map any keys to representations.")
+            return False
+
         # Download subtitles
         self.download_subtitles()
 
-        # Download the video to get segment count
+        # Download and decrypt video
         video_rep = self.get_representation_by_type("video")
         if video_rep:
+            video_key_info = key_mapping.get("video")
+            if not video_key_info:
+                self.error = "No key found for video representation"
+                return False
+                
+            console.log(f"[cyan]Using video key: [red]{video_key_info['kid']} [cyan]for representation [yellow]{video_key_info.get('representation_id')}")
+            
             video_downloader = MPD_Segments(tmp_folder=self.encrypted_dir, representation=video_rep, pssh=self.parser.pssh, custom_headers=custom_headers)
             encrypted_path = video_downloader.get_concat_path(self.encrypted_dir)
 
@@ -266,33 +278,28 @@ def download_and_decrypt(self, custom_headers=None, query_params=None, key=None)
                     self.current_downloader = None
                     self.current_download_type = None
 
-                # Try decryption with multiple keys if the first one fails
-                decrypted_path = os.path.join(self.decrypted_dir, f"video.{extension_output}")
-                result_path = None
-                
-                for i, key_info in enumerate(keys):
-                    KID = key_info['kid']
-                    KEY = key_info['key']
-                    console.log(f"[cyan]Trying video decryption with key: [red]{KID}")
-                    result_path = decrypt_with_mp4decrypt("Video", encrypted_path, KID, KEY, output_path=decrypted_path)
-                    
-                    if result_path:
-                        working_key = key_info      # Store the working key for audio decryption
-                        break
-                    else:
-                        console.log(f"[yellow]✗ Video decryption failed with key {i+1}")
+                # Decrypt video using the mapped key
+                decrypted_path = os.path.join(self.decrypted_dir, f"video.{EXTENSION_OUTPUT}")
+                result_path = decrypt_with_mp4decrypt("Video", encrypted_path, video_key_info['kid'], video_key_info['key'], output_path=decrypted_path)
 
                 if not result_path:
-                    self.error = "Video decryption failed with all available keys"
+                    self.error = f"Video decryption failed with key {video_key_info['kid']}"
                     return False
 
         else:
             self.error = "No video found"
             return False
 
-        # Now download audio with segment limiting
+        # Download and decrypt audio
         audio_rep = self.get_representation_by_type("audio")
         if audio_rep:
+            audio_key_info = key_mapping.get("audio")
+            if not audio_key_info:
+                self.error = "No key found for audio representation"
+                return False
+                
+            console.log(f"[cyan]Using audio key: [red]{audio_key_info['kid']} [cyan]for representation [yellow]{audio_key_info.get('representation_id')}")
+            
             audio_language = audio_rep.get('language', 'Unknown')
             audio_downloader = MPD_Segments(tmp_folder=self.encrypted_dir, representation=audio_rep, pssh=self.parser.pssh, limit_segments=video_segments_count if video_segments_count > 0 else None, custom_headers=custom_headers)
             encrypted_path = audio_downloader.get_concat_path(self.encrypted_dir)
@@ -325,32 +332,13 @@ def download_and_decrypt(self, custom_headers=None, query_params=None, key=None)
                     self.current_downloader = None
                     self.current_download_type = None
 
-                # Decrypt audio
-                decrypted_path = os.path.join(self.decrypted_dir, f"audio.{extension_output}")
-                result_path = decrypt_with_mp4decrypt(
-                    f"Audio {audio_language}", encrypted_path, working_key['kid'], working_key['key'], output_path=decrypted_path
-                )
+                # Decrypt audio using the mapped key
+                decrypted_path = os.path.join(self.decrypted_dir, f"audio.{EXTENSION_OUTPUT}")
+                result_path = decrypt_with_mp4decrypt(f"Audio {audio_language}", encrypted_path, audio_key_info['kid'], audio_key_info['key'], output_path=decrypted_path)
 
                 if not result_path:
-                    # If the video key doesn't work for audio, try other keys
-                    console.log("[yellow]Video key failed for audio, trying other keys...")
-                    
-                    for i, key_info in enumerate(keys):
-                        if key_info['kid'] == working_key['kid']:
-                            continue  # Skip the already tried key
-                            
-                        console.log(f"[cyan]Trying audio decryption with alternative key: {key_info['kid']}")
-                        result_path = decrypt_with_mp4decrypt(f"Audio {audio_language}", encrypted_path, key_info['kid'], key_info['key'], output_path=decrypted_path)
-                        
-                        if result_path:
-                            console.log("[green]Audio decryption successful with alternative key")
-                            break
-                        else:
-                            console.log("[yellow]Audio decryption failed with alternative key")
-                    
-                    if not result_path:
-                        self.error = "Audio decryption failed with all available keys"
-                        return False
+                    self.error = f"Audio decryption failed with key {audio_key_info['kid']}"
+                    return False
 
         else:
             self.error = "No audio found"
@@ -415,7 +403,7 @@ def download_segments(self, clear=False):
                     self.current_download_type = None
 
             # NO DECRYPTION: just copy/move to decrypted folder
-            decrypted_path = os.path.join(self.decrypted_dir, f"video.{extension_output}")
+            decrypted_path = os.path.join(self.decrypted_dir, f"video.{EXTENSION_OUTPUT}")
             if os.path.exists(encrypted_path) and not os.path.exists(decrypted_path):
                 shutil.copy2(encrypted_path, decrypted_path)
 
@@ -459,7 +447,7 @@ def download_segments(self, clear=False):
                     self.current_download_type = None
 
             # NO DECRYPTION: just copy/move to decrypted folder
-            decrypted_path = os.path.join(self.decrypted_dir, f"audio.{extension_output}")
+            decrypted_path = os.path.join(self.decrypted_dir, f"audio.{EXTENSION_OUTPUT}")
             if os.path.exists(encrypted_path) and not os.path.exists(decrypted_path):
                 shutil.copy2(encrypted_path, decrypted_path)
 
@@ -480,8 +468,8 @@ def finalize_output(self):
             return output_file
 
         # Definition of decrypted files
-        video_file = os.path.join(self.decrypted_dir, f"video.{extension_output}")
-        audio_file = os.path.join(self.decrypted_dir, f"audio.{extension_output}")
+        video_file = os.path.join(self.decrypted_dir, f"video.{EXTENSION_OUTPUT}")
+        audio_file = os.path.join(self.decrypted_dir, f"audio.{EXTENSION_OUTPUT}")
         output_file = self.original_output_path
 
         # Set the output file path for status tracking

StreamingCommunity/Lib/DASH/parser.py

@@ -222,6 +222,25 @@ def is_protected(self, element: etree._Element) -> bool:
 
         return False
 
+    def extract_default_kid(self, element: etree._Element) -> Optional[str]:
+        """Extract default_KID from ContentProtection element"""
+        for cp in self.ns.findall(element, 'mpd:ContentProtection'):
+            scheme_id = (cp.get('schemeIdUri') or '').lower()
+            
+            # Look for CENC protection with default_KID
+            if 'urn:mpeg:dash:mp4protection:2011' in scheme_id:
+                kid = cp.get('{urn:mpeg:cenc:2013}default_KID')
+                if not kid:
+                    kid = cp.get('cenc:default_KID')
+                if not kid:
+                    kid = cp.get('default_KID')
+                
+                if kid:
+                    # Remove dashes and return normalized KID
+                    return kid.replace('-', '').lower()
+        
+        return None
+    
     def extract_pssh(self, root: etree._Element) -> Optional[str]:
         """Extract PSSH (Protection System Specific Header)"""
         # Try Widevine first
@@ -446,12 +465,14 @@ def parse_adaptation_set(
         mime_type = adapt_set.get('mimeType', '')
         lang = adapt_set.get('lang', '')
         adapt_frame_rate = adapt_set.get('frameRate')
+        content_type = adapt_set.get('contentType', '')
 
         # Resolve base URL
         adapt_base = self.url_resolver.resolve_base_url(adapt_set, base_url)
 
-        # Check protection
+        # Check protection and extract default_KID
         adapt_protected = self.protection_handler.is_protected(adapt_set)
+        adapt_default_kid = self.protection_handler.extract_default_kid(adapt_set)
 
         # Get segment template
         adapt_seg_template = self.ns.find(adapt_set, 'mpd:SegmentTemplate')
@@ -469,6 +490,11 @@ def parse_adaptation_set(
                 rep['channels'] = self.metadata_extractor.get_audio_channels(rep_elem, adapt_set)
                 rep_protected = adapt_protected or self.protection_handler.is_protected(rep_elem)
                 rep['protected'] = bool(rep_protected)
+                rep_default_kid = self.protection_handler.extract_default_kid(rep_elem) or adapt_default_kid
+                rep['default_kid'] = rep_default_kid
+                if content_type:
+                    rep['type'] = content_type
+                
                 representations.append(rep)
 
         return representations