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