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_decrypt_text(self):
  87         password = intlist_to_bytes(self.key).decode()
  88         encrypted = base64.b64encode(
  89             intlist_to_bytes(self.iv[:8])
  90             + b'\x17\x15\x93\xab\x8d\x80V\xcdV\xe0\t\xcdo\xc2\xa5\xd8ksM\r\xe27N\xae'
  91         ).decode()
  92         decrypted = (aes_decrypt_text(encrypted, password, 16))
  93         self.assertEqual(decrypted, self.secret_msg)
  94
  95         password = intlist_to_bytes(self.key).decode()
  96         encrypted = base64.b64encode(
  97             intlist_to_bytes(self.iv[:8])
  98             + b'\x0b\xe6\xa4\xd9z\x0e\xb8\xb9\xd0\xd4i_\x85\x1d\x99\x98_\xe5\x80\xe7.\xbf\xa5\x83'
  99         ).decode()
 100         decrypted = (aes_decrypt_text(encrypted, password, 32))
 101         self.assertEqual(decrypted, self.secret_msg)
 102
 103     def test_ecb_encrypt(self):
 104         data = bytes_to_intlist(self.secret_msg)
 105         encrypted = intlist_to_bytes(aes_ecb_encrypt(data, self.key))
 106         self.assertEqual(
 107             encrypted,
 108             b'\xaa\x86]\x81\x97>\x02\x92\x9d\x1bR[[L/u\xd3&\xd1(h\xde{\x81\x94\xba\x02\xae\xbd\xa6\xd0:')
 109
 110     def test_ecb_decrypt(self):
 111         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:')
 112         decrypted = intlist_to_bytes(aes_ecb_decrypt(data, self.key, self.iv))
 113         self.assertEqual(decrypted.rstrip(b'\x08'), self.secret_msg)
 114
 115     def test_key_expansion(self):
 116         key = '4f6bdaa39e2f8cb07f5e722d9edef314'
 117
 118         self.assertEqual(key_expansion(bytes_to_intlist(bytearray.fromhex(key))), [
 119             0x4F, 0x6B, 0xDA, 0xA3, 0x9E, 0x2F, 0x8C, 0xB0, 0x7F, 0x5E, 0x72, 0x2D, 0x9E, 0xDE, 0xF3, 0x14,
 120             0x53, 0x66, 0x20, 0xA8, 0xCD, 0x49, 0xAC, 0x18, 0xB2, 0x17, 0xDE, 0x35, 0x2C, 0xC9, 0x2D, 0x21,
 121             0x8C, 0xBE, 0xDD, 0xD9, 0x41, 0xF7, 0x71, 0xC1, 0xF3, 0xE0, 0xAF, 0xF4, 0xDF, 0x29, 0x82, 0xD5,
 122             0x2D, 0xAD, 0xDE, 0x47, 0x6C, 0x5A, 0xAF, 0x86, 0x9F, 0xBA, 0x00, 0x72, 0x40, 0x93, 0x82, 0xA7,
 123             0xF9, 0xBE, 0x82, 0x4E, 0x95, 0xE4, 0x2D, 0xC8, 0x0A, 0x5E, 0x2D, 0xBA, 0x4A, 0xCD, 0xAF, 0x1D,
 124             0x54, 0xC7, 0x26, 0x98, 0xC1, 0x23, 0x0B, 0x50, 0xCB, 0x7D, 0x26, 0xEA, 0x81, 0xB0, 0x89, 0xF7,
 125             0x93, 0x60, 0x4E, 0x94, 0x52, 0x43, 0x45, 0xC4, 0x99, 0x3E, 0x63, 0x2E, 0x18, 0x8E, 0xEA, 0xD9,
 126             0xCA, 0xE7, 0x7B, 0x39, 0x98, 0xA4, 0x3E, 0xFD, 0x01, 0x9A, 0x5D, 0xD3, 0x19, 0x14, 0xB7, 0x0A,
 127             0xB0, 0x4E, 0x1C, 0xED, 0x28, 0xEA, 0x22, 0x10, 0x29, 0x70, 0x7F, 0xC3, 0x30, 0x64, 0xC8, 0xC9,
 128             0xE8, 0xA6, 0xC1, 0xE9, 0xC0, 0x4C, 0xE3, 0xF9, 0xE9, 0x3C, 0x9C, 0x3A, 0xD9, 0x58, 0x54, 0xF3,
 129             0xB4, 0x86, 0xCC, 0xDC, 0x74, 0xCA, 0x2F, 0x25, 0x9D, 0xF6, 0xB3, 0x1F, 0x44, 0xAE, 0xE7, 0xEC])
 130
 131     def test_pad_block(self):
 132         block = [0x21, 0xA0, 0x43, 0xFF]
 133
 134         self.assertEqual(pad_block(block, 'pkcs7'),
 135                          block + [0x0C, 0x0C, 0x0C, 0x0C, 0x0C, 0x0C, 0x0C, 0x0C, 0x0C, 0x0C, 0x0C, 0x0C])
 136
 137         self.assertEqual(pad_block(block, 'iso7816'),
 138                          block + [0x80, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00])
 139
 140         self.assertEqual(pad_block(block, 'whitespace'),
 141                          block + [0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20])
 142
 143         self.assertEqual(pad_block(block, 'zero'),
 144                          block + [0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00])
 145
 146         block = list(range(16))
 147         for mode in ('pkcs7', 'iso7816', 'whitespace', 'zero'):
 148             self.assertEqual(pad_block(block, mode), block, mode)
 149
 150
 151 if __name__ == '__main__':
 152     unittest.main()