[utils] Add improved jwt_encode function (#14071)

Also deprecates `jwt_encode_hs256`

Authored by: bashonly

yt_dlp/extractor/atvat.py

@@ -4,7 +4,7 @@ from .common import InfoExtractor
 from ..utils import (
     ExtractorError,
     float_or_none,
-    jwt_encode_hs256,
+    jwt_encode,
     try_get,
 )
 
@@ -83,11 +83,10 @@ class ATVAtIE(InfoExtractor):
             'nbf': int(not_before.timestamp()),
             'exp': int(expire.timestamp()),
         }
-        jwt_token = jwt_encode_hs256(payload, self._ENCRYPTION_KEY, headers={'kid': self._ACCESS_ID})
         videos = self._download_json(
             'https://vas-v4.p7s1video.net/4.0/getsources',
             content_id, 'Downloading videos JSON', query={
-                'token': jwt_token.decode('utf-8'),
+                'token': jwt_encode(payload, self._ENCRYPTION_KEY, headers={'kid': self._ACCESS_ID}),
             })
 
         video_id, videos_data = next(iter(videos['data'].items()))

yt_dlp/extractor/vrt.py

@@ -14,7 +14,7 @@ from ..utils import (
     get_element_html_by_class,
     int_or_none,
     jwt_decode_hs256,
-    jwt_encode_hs256,
+    jwt_encode,
     make_archive_id,
     merge_dicts,
     parse_age_limit,
@@ -98,9 +98,9 @@ class VRTBaseIE(InfoExtractor):
                 'Content-Type': 'application/json',
             }, data=json.dumps({
                 'identityToken': id_token or '',
-                'playerInfo': jwt_encode_hs256(player_info, self._JWT_SIGNING_KEY, headers={
+                'playerInfo': jwt_encode(player_info, self._JWT_SIGNING_KEY, headers={
                     'kid': self._JWT_KEY_ID,
-                }).decode(),
+                }),
             }, separators=(',', ':')).encode())['vrtPlayerToken']
 
         return self._download_json(

yt_dlp/utils/_deprecated.py

@@ -1,4 +1,8 @@
 """Deprecated - New code should avoid these"""
+import base64
+import hashlib
+import hmac
+import json
 import warnings
 
 from ..compat.compat_utils import passthrough_module
@@ -28,4 +32,18 @@ def intlist_to_bytes(xs):
     return struct.pack('%dB' % len(xs), *xs)
 
 
+def jwt_encode_hs256(payload_data, key, headers={}):
+    header_data = {
+        'alg': 'HS256',
+        'typ': 'JWT',
+    }
+    if headers:
+        header_data.update(headers)
+    header_b64 = base64.b64encode(json.dumps(header_data).encode())
+    payload_b64 = base64.b64encode(json.dumps(payload_data).encode())
+    h = hmac.new(key.encode(), header_b64 + b'.' + payload_b64, hashlib.sha256)
+    signature_b64 = base64.b64encode(h.digest())
+    return header_b64 + b'.' + payload_b64 + b'.' + signature_b64
+
+
 compiled_regex_type = type(re.compile(''))

yt_dlp/utils/_utils.py

@@ -4739,22 +4739,36 @@ def time_seconds(**kwargs):
     return time.time() + dt.timedelta(**kwargs).total_seconds()
 
 
-# create a JSON Web Signature (jws) with HS256 algorithm
-# the resulting format is in JWS Compact Serialization
 # implemented following JWT https://www.rfc-editor.org/rfc/rfc7519.html
 # implemented following JWS https://www.rfc-editor.org/rfc/rfc7515.html
-def jwt_encode_hs256(payload_data, key, headers={}):
+def jwt_encode(payload_data, key, *, alg='HS256', headers=None):
+    assert alg in ('HS256',), f'Unsupported algorithm "{alg}"'
+
+    def jwt_json_bytes(obj):
+        return json.dumps(obj, separators=(',', ':')).encode()
+
+    def jwt_b64encode(bytestring):
+        return base64.urlsafe_b64encode(bytestring).rstrip(b'=')
+
     header_data = {
-        'alg': 'HS256',
+        'alg': alg,
         'typ': 'JWT',
     }
     if headers:
-        header_data.update(headers)
-    header_b64 = base64.b64encode(json.dumps(header_data).encode())
-    payload_b64 = base64.b64encode(json.dumps(payload_data).encode())
+        # Allow re-ordering of keys if both 'alg' and 'typ' are present
+        if 'alg' in headers and 'typ' in headers:
+            header_data = headers
+        else:
+            header_data.update(headers)
+
+    header_b64 = jwt_b64encode(jwt_json_bytes(header_data))
+    payload_b64 = jwt_b64encode(jwt_json_bytes(payload_data))
+
+    # HS256 is the only algorithm currently supported
     h = hmac.new(key.encode(), header_b64 + b'.' + payload_b64, hashlib.sha256)
-    signature_b64 = base64.b64encode(h.digest())
-    return header_b64 + b'.' + payload_b64 + b'.' + signature_b64
+    signature_b64 = jwt_b64encode(h.digest())
+
+    return (header_b64 + b'.' + payload_b64 + b'.' + signature_b64).decode()
 
 
 # can be extended in future to verify the signature and parse header and return the algorithm used if it's not HS256