Repository: yt-dlp/yt-dlp
Pull request #1686 (2021-11-16 04:04:27)

AbemaTV batch 1/3: [AES] add ECB mode
## Please follow the guide below

- You will be asked some questions, please read them **carefully** and answer honestly
- Put an `x` into all the boxes [ ] relevant to your *pull request* (like that [x])
- Use *Preview* tab to see how your *pull request* will actually look like

---

### Before submitting a *pull request* make sure you have:
- [x] At least skimmed through [contributing guidelines](https://github.com/yt-dlp/yt-dlp/blob/master/CONTRIBUTING.md#developer-instructions) including [yt-dlp coding conventions](https://github.com/yt-dlp/yt-dlp/blob/master/CONTRIBUTING.md#yt-dlp-coding-conventions)
- [x] [Searched](https://github.com/yt-dlp/yt-dlp/search?q=is%3Apr&type=Issues) the bugtracker for similar pull requests
- [x] Checked the code with [flake8](https://pypi.python.org/pypi/flake8)

### In order to be accepted and merged into yt-dlp each piece of code must be in public domain or released under [Unlicense](http://unlicense.org/). Check one of the following options:
- [x] I am the original author of this code and I am willing to release it under [Unlicense](http://unlicense.org/)
- [ ] I am not the original author of this code but it is in public domain or released under [Unlicense](http://unlicense.org/) (provide reliable evidence)

### What is the purpose of your *pull request*?
- [ ] Bug fix
- [x] Improvement
- [ ] New extractor
- [ ] New feature

---

### Description of your *pull request* and other information

This PR adds support for ECB mode in aes.py, which is the simplest block cipher mode of operation.

This is mandatory for downloading from AbemaTV.


Patch:
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 @@ def aes_gcm_decrypt_and_verify_bytes(data, key, tag, nonce):
 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


Test patch:
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 @@
 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 @@
     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 @@ def test_decrypt_text(self):
         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()