From a04e005521ecf2eb0c4979e234ff0c4f23a3caa0 Mon Sep 17 00:00:00 2001
From: The Hatsune Daishi <nao20010128@gmail.com>
Date: Fri, 19 Nov 2021 10:54:10 +0900
Subject: [PATCH] [AES] Add ECB mode (#1686)

Needed for #1688
Authored by: nao20010128nao
---
 test/test_aes.py | 18 +++++++++++++++++-
 yt_dlp/aes.py    | 42 ++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 59 insertions(+), 1 deletion(-)

diff --git a/test/test_aes.py b/test/test_aes.py
index 46db59e57..5c9273f8a 100644
--- a/test/test_aes.py
+++ b/test/test_aes.py
@@ -10,6 +10,8 @@ sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 from yt_dlp.aes import (
     aes_decrypt,
     aes_encrypt,
+    aes_ecb_encrypt,
+    aes_ecb_decrypt,
     aes_cbc_decrypt,
     aes_cbc_decrypt_bytes,
     aes_cbc_encrypt,
@@ -17,7 +19,8 @@ from yt_dlp.aes import (
     aes_ctr_encrypt,
     aes_gcm_decrypt_and_verify,
     aes_gcm_decrypt_and_verify_bytes,
-    aes_decrypt_text
+    aes_decrypt_text,
+    BLOCK_SIZE_BYTES,
 )
 from yt_dlp.compat import compat_pycrypto_AES
 from yt_dlp.utils import bytes_to_intlist, intlist_to_bytes
@@ -94,6 +97,19 @@ class TestAES(unittest.TestCase):
         decrypted = (aes_decrypt_text(encrypted, password, 32))
         self.assertEqual(decrypted, self.secret_msg)
 
+    def test_ecb_encrypt(self):
+        data = bytes_to_intlist(self.secret_msg)
+        data += [0x08] * (BLOCK_SIZE_BYTES - len(data) % BLOCK_SIZE_BYTES)
+        encrypted = intlist_to_bytes(aes_ecb_encrypt(data, self.key, self.iv))
+        self.assertEqual(
+            encrypted,
+            b'\xaa\x86]\x81\x97>\x02\x92\x9d\x1bR[[L/u\xd3&\xd1(h\xde{\x81\x94\xba\x02\xae\xbd\xa6\xd0:')
+
+    def test_ecb_decrypt(self):
+        data = bytes_to_intlist(b'\xaa\x86]\x81\x97>\x02\x92\x9d\x1bR[[L/u\xd3&\xd1(h\xde{\x81\x94\xba\x02\xae\xbd\xa6\xd0:')
+        decrypted = intlist_to_bytes(aes_ecb_decrypt(data, self.key, self.iv))
+        self.assertEqual(decrypted.rstrip(b'\x08'), self.secret_msg)
+
 
 if __name__ == '__main__':
     unittest.main()
diff --git a/yt_dlp/aes.py b/yt_dlp/aes.py
index 60cdeb74e..8503e3dfd 100644
--- a/yt_dlp/aes.py
+++ b/yt_dlp/aes.py
@@ -28,6 +28,48 @@ else:
 BLOCK_SIZE_BYTES = 16
 
 
+def aes_ecb_encrypt(data, key, iv=None):
+    """
+    Encrypt with aes in ECB mode
+
+    @param {int[]} data        cleartext
+    @param {int[]} key         16/24/32-Byte cipher key
+    @param {int[]} iv          Unused for this mode
+    @returns {int[]}           encrypted data
+    """
+    expanded_key = key_expansion(key)
+    block_count = int(ceil(float(len(data)) / BLOCK_SIZE_BYTES))
+
+    encrypted_data = []
+    for i in range(block_count):
+        block = data[i * BLOCK_SIZE_BYTES: (i + 1) * BLOCK_SIZE_BYTES]
+        encrypted_data += aes_encrypt(block, expanded_key)
+    encrypted_data = encrypted_data[:len(data)]
+
+    return encrypted_data
+
+
+def aes_ecb_decrypt(data, key, iv=None):
+    """
+    Decrypt with aes in ECB mode
+
+    @param {int[]} data        cleartext
+    @param {int[]} key         16/24/32-Byte cipher key
+    @param {int[]} iv          Unused for this mode
+    @returns {int[]}           decrypted data
+    """
+    expanded_key = key_expansion(key)
+    block_count = int(ceil(float(len(data)) / BLOCK_SIZE_BYTES))
+
+    encrypted_data = []
+    for i in range(block_count):
+        block = data[i * BLOCK_SIZE_BYTES: (i + 1) * BLOCK_SIZE_BYTES]
+        encrypted_data += aes_decrypt(block, expanded_key)
+    encrypted_data = encrypted_data[:len(data)]
+
+    return encrypted_data
+
+
 def aes_ctr_decrypt(data, key, iv):
     """
     Decrypt with aes in counter mode