From 4b985d7d7115d9a73c7300e2ad93b8e8c8674ca4 Mon Sep 17 00:00:00 2001
From: Sasuke-Duck UwU <95662926+SASUKE-DUCK@users.noreply.github.com>
Date: Wed, 15 Nov 2023 02:57:18 -0500
Subject: [PATCH] Add files via upload

Added ism pssh support and an ism parser just to get the information for video and audio codecs.
---
 cdm/wks.py | 75 +++++++++++++++++++++++++++++++++++++++++++++++++++++-
 ism.py     | 36 ++++++++++++++++++++++++++
 2 files changed, 110 insertions(+), 1 deletion(-)
 create mode 100644 ism.py

diff --git a/cdm/wks.py b/cdm/wks.py
index b6eb9e4..1d27832 100644
--- a/cdm/wks.py
+++ b/cdm/wks.py
@@ -14,7 +14,10 @@ import requests
 from base64 import b64encode
 from google.protobuf.message import DecodeError
 from google.protobuf import text_format
-
+import xmltodict
+import base64
+import uuid
+import requests
 from Cryptodome.Random import get_random_bytes
 from Cryptodome.Random import random
 from Cryptodome.Cipher import PKCS1_OAEP, AES
@@ -783,3 +786,73 @@ class DataExtractor_DSNP:
             matches = [(match[0], re.search(r'base64,(.*)', match[1]).group(1)) for match in re.findall(r'KEYFORMAT="urn:uuid:edef8ba9-79d6-4ace-a3c8-27dcd51d21ed",KEYFORMATVERSIONS="[^"]+",CHARACTERISTICS="([^"]+)",URI="([^"]+)"', self.content)]
             return matches
         return []
+
+def parse_manifest_ism(manifest_url):
+    r = requests.get(manifest_url, headers={'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) '
+                                                           'AppleWebKit/537.36 (KHTML, like Gecko) '
+                                                           'Chrome/72.0.3626.121 Safari/537.36'})
+
+    if r.status_code != 200:
+        raise Exception(r.text)
+
+    ism = xmltodict.parse(r.text)
+
+    pssh = ism['SmoothStreamingMedia']['Protection']['ProtectionHeader']['#text']
+
+    pr_pssh_dec = base64.b64decode(pssh).decode('utf16')
+    pr_pssh_dec = pr_pssh_dec[pr_pssh_dec.index('<'):]
+    pr_pssh_xml = xmltodict.parse(pr_pssh_dec)
+    kid_hex = base64.b64decode(pr_pssh_xml['WRMHEADER']['DATA']['KID']).hex()
+    
+    kid = uuid.UUID(kid_hex).bytes_le.hex()
+
+    stream_indices = ism['SmoothStreamingMedia']['StreamIndex']
+
+    # List to store information for each stream
+    stream_info_list = []
+
+    # Iterate over each StreamIndex (as it might be a list)
+    for stream_info in stream_indices if isinstance(stream_indices, list) else [stream_indices]:
+        type_info = stream_info['@Type']
+
+        if type_info in {'video', 'audio'}:
+            # Handle the case where there can be multiple QualityLevel elements
+            quality_levels = stream_info.get('QualityLevel', [])
+
+            if not isinstance(quality_levels, list):
+                quality_levels = [quality_levels]
+
+            for quality_level in quality_levels:
+                codec = quality_level.get('@FourCC', 'N/A')
+                bitrate = quality_level.get('@Bitrate', 'N/A')
+                
+                # Additional attributes for video streams
+                if type_info == 'video':
+                    max_width = quality_level.get('@MaxWidth', 'N/A')
+                    max_height = quality_level.get('@MaxHeight', 'N/A')
+                    resolution = f"{max_width}x{max_height}"
+                else:
+                    resolution = 'N/A'
+                
+                # Additional attributes for audio streams
+                language = stream_info.get('@Language', 'N/A')
+                track_id = stream_info.get('@AudioTrackId', 'N/A') if type_info == 'audio' else None
+
+                stream_info_list.append({
+                    'type': type_info,
+                    'codec': codec,
+                    'bitrate': bitrate,
+                    'resolution': resolution,
+                    'language': language,
+                    'track_id': track_id
+                })
+
+    # PSSH encoding logic in ism
+    array_of_bytes = bytearray(b'\x00\x00\x002pssh\x00\x00\x00\x00')
+    array_of_bytes.extend(bytes.fromhex("edef8ba979d64acea3c827dcd51d21ed"))
+    array_of_bytes.extend(b'\x00\x00\x00\x12\x12\x10')
+    array_of_bytes.extend(bytes.fromhex(str(kid).replace("-", "")))
+
+    encoded_string = base64.b64encode(bytes.fromhex(array_of_bytes.hex())).decode("utf-8")
+
+    return kid, stream_info_list, encoded_string
\ No newline at end of file
diff --git a/ism.py b/ism.py
new file mode 100644
index 0000000..93b3b12
--- /dev/null
+++ b/ism.py
@@ -0,0 +1,36 @@
+import argparse
+from cdm.wks import parse_manifest_ism
+
+def main():
+    # Create an ArgumentParser object and add the 'urls' argument
+    parser = argparse.ArgumentParser(description='Script for parsing Smooth Streaming manifest URLs.')
+    parser.add_argument('urls',
+                        help='The URLs to parse. You may need to wrap the URLs in double quotes if you have issues.',
+                        nargs='+')
+
+    # Parse the arguments
+    args = parser.parse_args()
+
+    # Iterate over the provided URLs
+    for manifest_link in args.urls:
+        kid, stream_info_list, encoded_string = parse_manifest_ism(manifest_link)
+
+        # Print information for each stream
+        for stream_info in stream_info_list:
+            type_info = stream_info['type']
+            codec = stream_info['codec']
+            bitrate = stream_info['bitrate']
+            resolution = stream_info['resolution']
+
+            if type_info == 'video':
+                print(f'[INFO] VIDEO - Codec: {codec}, Resolution: {resolution}, Bitrate: {bitrate}')
+            elif type_info == 'audio':
+                language = stream_info['language']
+                track_id = stream_info['track_id']
+                print(f'[INFO] AUDIO - Codec: {codec}, Bitrate: {bitrate}, Language: {language}, Track ID: {track_id}')
+
+        # Print PSSH information
+        print('\n[INFO] PSSH:', encoded_string)
+
+if __name__ == "__main__":
+    main()
\ No newline at end of file