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