test/test_aes.py

   1 #!/usr/bin/env python3
   2
   3 # Allow direct execution
   4 import os
   5 import sys
   6 import unittest
   7
   8 sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
   9
  10
  11 import base64
  12
  13 from yt_dlp.aes import (
  14     BLOCK_SIZE_BYTES,
  15     aes_cbc_decrypt,
  16     aes_cbc_decrypt_bytes,
  17     aes_cbc_encrypt,
  18     aes_ctr_decrypt,
  19     aes_ctr_encrypt,
  20     aes_decrypt,
  21     aes_decrypt_text,
  22     aes_ecb_decrypt,
  23     aes_ecb_encrypt,
  24     aes_encrypt,
  25     aes_gcm_decrypt_and_verify,
  26     aes_gcm_decrypt_and_verify_bytes,
  27     key_expansion,
  28     pad_block,
  29 )
  30 from yt_dlp.dependencies import Cryptodome_AES
  31 from yt_dlp.utils import bytes_to_intlist, intlist_to_bytes
  32
  33 # the encrypted data can be generate with 'devscripts/generate_aes_testdata.py'
  34
  35
  36 class TestAES(unittest.TestCase):
  37     def setUp(self):
  38         self.key = self.iv = [0x20, 0x15] + 14 * [0]
  39         self.secret_msg = b'Secret message goes here'
  40
  41     def test_encrypt(self):
  42         msg = b'message'
  43         key = list(range(16))
  44         encrypted = aes_encrypt(bytes_to_intlist(msg), key)
  45         decrypted = intlist_to_bytes(aes_decrypt(encrypted, key))
  46         self.assertEqual(decrypted, msg)
  47
  48     def test_cbc_decrypt(self):
  49         data = b'\x97\x92+\xe5\x0b\xc3\x18\x91ky9m&\xb3\xb5@\xe6\x27\xc2\x96.\xc8u\x88\xab9-[\x9e|\xf1\xcd'
  50         decrypted = intlist_to_bytes(aes_cbc_decrypt(bytes_to_intlist(data), self.key, self.iv))
  51         self.assertEqual(decrypted.rstrip(b'\x08'), self.secret_msg)
  52         if Cryptodome_AES:
  53             decrypted = aes_cbc_decrypt_bytes(data, intlist_to_bytes(self.key), intlist_to_bytes(self.iv))
  54             self.assertEqual(decrypted.rstrip(b'\x08'), self.secret_msg)
  55
  56     def test_cbc_encrypt(self):
  57         data = bytes_to_intlist(self.secret_msg)
  58         encrypted = intlist_to_bytes(aes_cbc_encrypt(data, self.key, self.iv))
  59         self.assertEqual(
  60             encrypted,
  61             b'\x97\x92+\xe5\x0b\xc3\x18\x91ky9m&\xb3\xb5@\xe6\'\xc2\x96.\xc8u\x88\xab9-[\x9e|\xf1\xcd')
  62
  63     def test_ctr_decrypt(self):
  64         data = bytes_to_intlist(b'\x03\xc7\xdd\xd4\x8e\xb3\xbc\x1a*O\xdc1\x12+8Aio\xd1z\xb5#\xaf\x08')
  65         decrypted = intlist_to_bytes(aes_ctr_decrypt(data, self.key, self.iv))
  66         self.assertEqual(decrypted.rstrip(b'\x08'), self.secret_msg)
  67
  68     def test_ctr_encrypt(self):
  69         data = bytes_to_intlist(self.secret_msg)
  70         encrypted = intlist_to_bytes(aes_ctr_encrypt(data, self.key, self.iv))
  71         self.assertEqual(
  72             encrypted,
  73             b'\x03\xc7\xdd\xd4\x8e\xb3\xbc\x1a*O\xdc1\x12+8Aio\xd1z\xb5#\xaf\x08')
  74
  75     def test_gcm_decrypt(self):
  76         data = b'\x159Y\xcf5eud\x90\x9c\x85&]\x14\x1d\x0f.\x08\xb4T\xe4/\x17\xbd'
  77         authentication_tag = b'\xe8&I\x80rI\x07\x9d}YWuU@:e'
  78
  79         decrypted = intlist_to_bytes(aes_gcm_decrypt_and_verify(
  80             bytes_to_intlist(data), self.key, bytes_to_intlist(authentication_tag), self.iv[:12]))
  81         self.assertEqual(decrypted.rstrip(b'\x08'), self.secret_msg)
  82         if Cryptodome_AES:
  83             decrypted = aes_gcm_decrypt_and_verify_bytes(
  84                 data, intlist_to_bytes(self.key), authentication_tag, intlist_to_bytes(self.iv[:12]))
  85             self.assertEqual(decrypted.rstrip(b'\x08'), self.secret_msg)
  86
  87     def test_decrypt_text(self):
  88         password = intlist_to_bytes(self.key).decode()
  89         encrypted = base64.b64encode(
  90             intlist_to_bytes(self.iv[:8])
  91             + b'\x17\x15\x93\xab\x8d\x80V\xcdV\xe0\t\xcdo\xc2\xa5\xd8ksM\r\xe27N\xae'
  92         ).decode()
  93         decrypted = (aes_decrypt_text(encrypted, password, 16))
  94         self.assertEqual(decrypted, self.secret_msg)
  95
  96         password = intlist_to_bytes(self.key).decode()
  97         encrypted = base64.b64encode(
  98             intlist_to_bytes(self.iv[:8])
  99             + b'\x0b\xe6\xa4\xd9z\x0e\xb8\xb9\xd0\xd4i_\x85\x1d\x99\x98_\xe5\x80\xe7.\xbf\xa5\x83'
 100         ).decode()
 101         decrypted = (aes_decrypt_text(encrypted, password, 32))
 102         self.assertEqual(decrypted, self.secret_msg)
 103
 104     def test_ecb_encrypt(self):
 105         data = bytes_to_intlist(self.secret_msg)
 106         data += [0x08] * (BLOCK_SIZE_BYTES - len(data) % BLOCK_SIZE_BYTES)
 107         encrypted = intlist_to_bytes(aes_ecb_encrypt(data, self.key, self.iv))
 108         self.assertEqual(
 109             encrypted,
 110             b'\xaa\x86]\x81\x97>\x02\x92\x9d\x1bR[[L/u\xd3&\xd1(h\xde{\x81\x94\xba\x02\xae\xbd\xa6\xd0:')
 111
 112     def test_ecb_decrypt(self):
 113         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:')
 114         decrypted = intlist_to_bytes(aes_ecb_decrypt(data, self.key, self.iv))
 115         self.assertEqual(decrypted.rstrip(b'\x08'), self.secret_msg)
 116
 117     def test_key_expansion(self):
 118         key = '4f6bdaa39e2f8cb07f5e722d9edef314'
 119
 120         self.assertEqual(key_expansion(bytes_to_intlist(bytearray.fromhex(key))), [
 121             0x4F, 0x6B, 0xDA, 0xA3, 0x9E, 0x2F, 0x8C, 0xB0, 0x7F, 0x5E, 0x72, 0x2D, 0x9E, 0xDE, 0xF3, 0x14,
 122             0x53, 0x66, 0x20, 0xA8, 0xCD, 0x49, 0xAC, 0x18, 0xB2, 0x17, 0xDE, 0x35, 0x2C, 0xC9, 0x2D, 0x21,
 123             0x8C, 0xBE, 0xDD, 0xD9, 0x41, 0xF7, 0x71, 0xC1, 0xF3, 0xE0, 0xAF, 0xF4, 0xDF, 0x29, 0x82, 0xD5,
 124             0x2D, 0xAD, 0xDE, 0x47, 0x6C, 0x5A, 0xAF, 0x86, 0x9F, 0xBA, 0x00, 0x72, 0x40, 0x93, 0x82, 0xA7,
 125             0xF9, 0xBE, 0x82, 0x4E, 0x95, 0xE4, 0x2D, 0xC8, 0x0A, 0x5E, 0x2D, 0xBA, 0x4A, 0xCD, 0xAF, 0x1D,
 126             0x54, 0xC7, 0x26, 0x98, 0xC1, 0x23, 0x0B, 0x50, 0xCB, 0x7D, 0x26, 0xEA, 0x81, 0xB0, 0x89, 0xF7,
 127             0x93, 0x60, 0x4E, 0x94, 0x52, 0x43, 0x45, 0xC4, 0x99, 0x3E, 0x63, 0x2E, 0x18, 0x8E, 0xEA, 0xD9,
 128             0xCA, 0xE7, 0x7B, 0x39, 0x98, 0xA4, 0x3E, 0xFD, 0x01, 0x9A, 0x5D, 0xD3, 0x19, 0x14, 0xB7, 0x0A,
 129             0xB0, 0x4E, 0x1C, 0xED, 0x28, 0xEA, 0x22, 0x10, 0x29, 0x70, 0x7F, 0xC3, 0x30, 0x64, 0xC8, 0xC9,
 130             0xE8, 0xA6, 0xC1, 0xE9, 0xC0, 0x4C, 0xE3, 0xF9, 0xE9, 0x3C, 0x9C, 0x3A, 0xD9, 0x58, 0x54, 0xF3,
 131             0xB4, 0x86, 0xCC, 0xDC, 0x74, 0xCA, 0x2F, 0x25, 0x9D, 0xF6, 0xB3, 0x1F, 0x44, 0xAE, 0xE7, 0xEC])
 132
 133     def test_pad_block(self):
 134         block = [0x21, 0xA0, 0x43, 0xFF]
 135
 136         self.assertEqual(pad_block(block, 'pkcs7'),
 137                          block + [0x0C, 0x0C, 0x0C, 0x0C, 0x0C, 0x0C, 0x0C, 0x0C, 0x0C, 0x0C, 0x0C, 0x0C])
 138
 139         self.assertEqual(pad_block(block, 'iso7816'),
 140                          block + [0x80, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00])
 141
 142         self.assertEqual(pad_block(block, 'whitespace'),
 143                          block + [0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20])
 144
 145         self.assertEqual(pad_block(block, 'zero'),
 146                          block + [0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00])
 147
 148         block = list(range(16))
 149         for mode in ('pkcs7', 'iso7816', 'whitespace', 'zero'):
 150             self.assertEqual(pad_block(block, mode), block, mode)
 151
 152
 153 if __name__ == '__main__':
 154     unittest.main()