Merge remote-tracking branch 'origin/master'
[ytdl] / youtube_dl / aes.py
1 __all__ = ['aes_encrypt', 'key_expansion', 'aes_ctr_decrypt', 'aes_cbc_decrypt', 'aes_decrypt_text']
2
3 import base64
4 from math import ceil
5
6 from .utils import bytes_to_intlist, intlist_to_bytes
7
8 BLOCK_SIZE_BYTES = 16
9
10 def aes_ctr_decrypt(data, key, counter):
11     """
12     Decrypt with aes in counter mode
13     
14     @param {int[]} data        cipher
15     @param {int[]} key         16/24/32-Byte cipher key
16     @param {instance} counter  Instance whose next_value function (@returns {int[]}  16-Byte block)
17                                returns the next counter block
18     @returns {int[]}           decrypted data
19     """
20     expanded_key = key_expansion(key)
21     block_count = int(ceil(float(len(data)) / BLOCK_SIZE_BYTES))
22     
23     decrypted_data=[]
24     for i in range(block_count):
25         counter_block = counter.next_value()
26         block = data[i*BLOCK_SIZE_BYTES : (i+1)*BLOCK_SIZE_BYTES]
27         block += [0]*(BLOCK_SIZE_BYTES - len(block))
28         
29         cipher_counter_block = aes_encrypt(counter_block, expanded_key)
30         decrypted_data += xor(block, cipher_counter_block)
31     decrypted_data = decrypted_data[:len(data)]
32     
33     return decrypted_data
34
35 def aes_cbc_decrypt(data, key, iv):
36     """
37     Decrypt with aes in CBC mode
38     
39     @param {int[]} data        cipher
40     @param {int[]} key         16/24/32-Byte cipher key
41     @param {int[]} iv          16-Byte IV
42     @returns {int[]}           decrypted data
43     """
44     expanded_key = key_expansion(key)
45     block_count = int(ceil(float(len(data)) / BLOCK_SIZE_BYTES))
46     
47     decrypted_data=[]
48     previous_cipher_block = iv
49     for i in range(block_count):
50         block = data[i*BLOCK_SIZE_BYTES : (i+1)*BLOCK_SIZE_BYTES]
51         block += [0]*(BLOCK_SIZE_BYTES - len(block))
52         
53         decrypted_block = aes_decrypt(block, expanded_key)
54         decrypted_data += xor(decrypted_block, previous_cipher_block)
55         previous_cipher_block = block
56     decrypted_data = decrypted_data[:len(data)]
57     
58     return decrypted_data
59
60 def key_expansion(data):
61     """
62     Generate key schedule
63     
64     @param {int[]} data  16/24/32-Byte cipher key
65     @returns {int[]}     176/208/240-Byte expanded key 
66     """
67     data = data[:] # copy
68     rcon_iteration = 1
69     key_size_bytes = len(data)
70     expanded_key_size_bytes = (key_size_bytes // 4 + 7) * BLOCK_SIZE_BYTES
71     
72     while len(data) < expanded_key_size_bytes:
73         temp = data[-4:]
74         temp = key_schedule_core(temp, rcon_iteration)
75         rcon_iteration += 1
76         data += xor(temp, data[-key_size_bytes : 4-key_size_bytes])
77         
78         for _ in range(3):
79             temp = data[-4:]
80             data += xor(temp, data[-key_size_bytes : 4-key_size_bytes])
81         
82         if key_size_bytes == 32:
83             temp = data[-4:]
84             temp = sub_bytes(temp)
85             data += xor(temp, data[-key_size_bytes : 4-key_size_bytes])
86         
87         for _ in range(3 if key_size_bytes == 32  else 2 if key_size_bytes == 24 else 0):
88             temp = data[-4:]
89             data += xor(temp, data[-key_size_bytes : 4-key_size_bytes])
90     data = data[:expanded_key_size_bytes]
91     
92     return data
93
94 def aes_encrypt(data, expanded_key):
95     """
96     Encrypt one block with aes
97     
98     @param {int[]} data          16-Byte state
99     @param {int[]} expanded_key  176/208/240-Byte expanded key 
100     @returns {int[]}             16-Byte cipher
101     """
102     rounds = len(expanded_key) // BLOCK_SIZE_BYTES - 1
103
104     data = xor(data, expanded_key[:BLOCK_SIZE_BYTES])
105     for i in range(1, rounds+1):
106         data = sub_bytes(data)
107         data = shift_rows(data)
108         if i != rounds:
109             data = mix_columns(data)
110         data = xor(data, expanded_key[i*BLOCK_SIZE_BYTES : (i+1)*BLOCK_SIZE_BYTES])
111
112     return data
113
114 def aes_decrypt(data, expanded_key):
115     """
116     Decrypt one block with aes
117     
118     @param {int[]} data          16-Byte cipher
119     @param {int[]} expanded_key  176/208/240-Byte expanded key
120     @returns {int[]}             16-Byte state
121     """
122     rounds = len(expanded_key) // BLOCK_SIZE_BYTES - 1
123     
124     for i in range(rounds, 0, -1):
125         data = xor(data, expanded_key[i*BLOCK_SIZE_BYTES : (i+1)*BLOCK_SIZE_BYTES])
126         if i != rounds:
127             data = mix_columns_inv(data)
128         data = shift_rows_inv(data)
129         data = sub_bytes_inv(data)
130     data = xor(data, expanded_key[:BLOCK_SIZE_BYTES])
131     
132     return data
133
134 def aes_decrypt_text(data, password, key_size_bytes):
135     """
136     Decrypt text
137     - The first 8 Bytes of decoded 'data' are the 8 high Bytes of the counter
138     - The cipher key is retrieved by encrypting the first 16 Byte of 'password'
139       with the first 'key_size_bytes' Bytes from 'password' (if necessary filled with 0's)
140     - Mode of operation is 'counter'
141     
142     @param {str} data                    Base64 encoded string
143     @param {str,unicode} password        Password (will be encoded with utf-8)
144     @param {int} key_size_bytes          Possible values: 16 for 128-Bit, 24 for 192-Bit or 32 for 256-Bit
145     @returns {str}                       Decrypted data
146     """
147     NONCE_LENGTH_BYTES = 8
148     
149     data = bytes_to_intlist(base64.b64decode(data))
150     password = bytes_to_intlist(password.encode('utf-8'))
151     
152     key = password[:key_size_bytes] + [0]*(key_size_bytes - len(password))
153     key = aes_encrypt(key[:BLOCK_SIZE_BYTES], key_expansion(key)) * (key_size_bytes // BLOCK_SIZE_BYTES)
154     
155     nonce = data[:NONCE_LENGTH_BYTES]
156     cipher = data[NONCE_LENGTH_BYTES:]
157     
158     class Counter:
159         __value = nonce + [0]*(BLOCK_SIZE_BYTES - NONCE_LENGTH_BYTES)
160         def next_value(self):
161             temp = self.__value
162             self.__value = inc(self.__value)
163             return temp
164     
165     decrypted_data = aes_ctr_decrypt(cipher, key, Counter())
166     plaintext = intlist_to_bytes(decrypted_data)
167     
168     return plaintext
169
170 RCON = (0x8d, 0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, 0x1b, 0x36)
171 SBOX = (0x63, 0x7C, 0x77, 0x7B, 0xF2, 0x6B, 0x6F, 0xC5, 0x30, 0x01, 0x67, 0x2B, 0xFE, 0xD7, 0xAB, 0x76,
172         0xCA, 0x82, 0xC9, 0x7D, 0xFA, 0x59, 0x47, 0xF0, 0xAD, 0xD4, 0xA2, 0xAF, 0x9C, 0xA4, 0x72, 0xC0,
173         0xB7, 0xFD, 0x93, 0x26, 0x36, 0x3F, 0xF7, 0xCC, 0x34, 0xA5, 0xE5, 0xF1, 0x71, 0xD8, 0x31, 0x15,
174         0x04, 0xC7, 0x23, 0xC3, 0x18, 0x96, 0x05, 0x9A, 0x07, 0x12, 0x80, 0xE2, 0xEB, 0x27, 0xB2, 0x75,
175         0x09, 0x83, 0x2C, 0x1A, 0x1B, 0x6E, 0x5A, 0xA0, 0x52, 0x3B, 0xD6, 0xB3, 0x29, 0xE3, 0x2F, 0x84,
176         0x53, 0xD1, 0x00, 0xED, 0x20, 0xFC, 0xB1, 0x5B, 0x6A, 0xCB, 0xBE, 0x39, 0x4A, 0x4C, 0x58, 0xCF,
177         0xD0, 0xEF, 0xAA, 0xFB, 0x43, 0x4D, 0x33, 0x85, 0x45, 0xF9, 0x02, 0x7F, 0x50, 0x3C, 0x9F, 0xA8,
178         0x51, 0xA3, 0x40, 0x8F, 0x92, 0x9D, 0x38, 0xF5, 0xBC, 0xB6, 0xDA, 0x21, 0x10, 0xFF, 0xF3, 0xD2,
179         0xCD, 0x0C, 0x13, 0xEC, 0x5F, 0x97, 0x44, 0x17, 0xC4, 0xA7, 0x7E, 0x3D, 0x64, 0x5D, 0x19, 0x73,
180         0x60, 0x81, 0x4F, 0xDC, 0x22, 0x2A, 0x90, 0x88, 0x46, 0xEE, 0xB8, 0x14, 0xDE, 0x5E, 0x0B, 0xDB,
181         0xE0, 0x32, 0x3A, 0x0A, 0x49, 0x06, 0x24, 0x5C, 0xC2, 0xD3, 0xAC, 0x62, 0x91, 0x95, 0xE4, 0x79,
182         0xE7, 0xC8, 0x37, 0x6D, 0x8D, 0xD5, 0x4E, 0xA9, 0x6C, 0x56, 0xF4, 0xEA, 0x65, 0x7A, 0xAE, 0x08,
183         0xBA, 0x78, 0x25, 0x2E, 0x1C, 0xA6, 0xB4, 0xC6, 0xE8, 0xDD, 0x74, 0x1F, 0x4B, 0xBD, 0x8B, 0x8A,
184         0x70, 0x3E, 0xB5, 0x66, 0x48, 0x03, 0xF6, 0x0E, 0x61, 0x35, 0x57, 0xB9, 0x86, 0xC1, 0x1D, 0x9E,
185         0xE1, 0xF8, 0x98, 0x11, 0x69, 0xD9, 0x8E, 0x94, 0x9B, 0x1E, 0x87, 0xE9, 0xCE, 0x55, 0x28, 0xDF,
186         0x8C, 0xA1, 0x89, 0x0D, 0xBF, 0xE6, 0x42, 0x68, 0x41, 0x99, 0x2D, 0x0F, 0xB0, 0x54, 0xBB, 0x16)
187 SBOX_INV = (0x52, 0x09, 0x6a, 0xd5, 0x30, 0x36, 0xa5, 0x38, 0xbf, 0x40, 0xa3, 0x9e, 0x81, 0xf3, 0xd7, 0xfb,
188             0x7c, 0xe3, 0x39, 0x82, 0x9b, 0x2f, 0xff, 0x87, 0x34, 0x8e, 0x43, 0x44, 0xc4, 0xde, 0xe9, 0xcb,
189             0x54, 0x7b, 0x94, 0x32, 0xa6, 0xc2, 0x23, 0x3d, 0xee, 0x4c, 0x95, 0x0b, 0x42, 0xfa, 0xc3, 0x4e,
190             0x08, 0x2e, 0xa1, 0x66, 0x28, 0xd9, 0x24, 0xb2, 0x76, 0x5b, 0xa2, 0x49, 0x6d, 0x8b, 0xd1, 0x25,
191             0x72, 0xf8, 0xf6, 0x64, 0x86, 0x68, 0x98, 0x16, 0xd4, 0xa4, 0x5c, 0xcc, 0x5d, 0x65, 0xb6, 0x92,
192             0x6c, 0x70, 0x48, 0x50, 0xfd, 0xed, 0xb9, 0xda, 0x5e, 0x15, 0x46, 0x57, 0xa7, 0x8d, 0x9d, 0x84,
193             0x90, 0xd8, 0xab, 0x00, 0x8c, 0xbc, 0xd3, 0x0a, 0xf7, 0xe4, 0x58, 0x05, 0xb8, 0xb3, 0x45, 0x06,
194             0xd0, 0x2c, 0x1e, 0x8f, 0xca, 0x3f, 0x0f, 0x02, 0xc1, 0xaf, 0xbd, 0x03, 0x01, 0x13, 0x8a, 0x6b,
195             0x3a, 0x91, 0x11, 0x41, 0x4f, 0x67, 0xdc, 0xea, 0x97, 0xf2, 0xcf, 0xce, 0xf0, 0xb4, 0xe6, 0x73,
196             0x96, 0xac, 0x74, 0x22, 0xe7, 0xad, 0x35, 0x85, 0xe2, 0xf9, 0x37, 0xe8, 0x1c, 0x75, 0xdf, 0x6e,
197             0x47, 0xf1, 0x1a, 0x71, 0x1d, 0x29, 0xc5, 0x89, 0x6f, 0xb7, 0x62, 0x0e, 0xaa, 0x18, 0xbe, 0x1b,
198             0xfc, 0x56, 0x3e, 0x4b, 0xc6, 0xd2, 0x79, 0x20, 0x9a, 0xdb, 0xc0, 0xfe, 0x78, 0xcd, 0x5a, 0xf4,
199             0x1f, 0xdd, 0xa8, 0x33, 0x88, 0x07, 0xc7, 0x31, 0xb1, 0x12, 0x10, 0x59, 0x27, 0x80, 0xec, 0x5f,
200             0x60, 0x51, 0x7f, 0xa9, 0x19, 0xb5, 0x4a, 0x0d, 0x2d, 0xe5, 0x7a, 0x9f, 0x93, 0xc9, 0x9c, 0xef,
201             0xa0, 0xe0, 0x3b, 0x4d, 0xae, 0x2a, 0xf5, 0xb0, 0xc8, 0xeb, 0xbb, 0x3c, 0x83, 0x53, 0x99, 0x61,
202             0x17, 0x2b, 0x04, 0x7e, 0xba, 0x77, 0xd6, 0x26, 0xe1, 0x69, 0x14, 0x63, 0x55, 0x21, 0x0c, 0x7d)
203 MIX_COLUMN_MATRIX = ((0x2,0x3,0x1,0x1),
204                      (0x1,0x2,0x3,0x1),
205                      (0x1,0x1,0x2,0x3),
206                      (0x3,0x1,0x1,0x2))
207 MIX_COLUMN_MATRIX_INV = ((0xE,0xB,0xD,0x9),
208                          (0x9,0xE,0xB,0xD),
209                          (0xD,0x9,0xE,0xB),
210                          (0xB,0xD,0x9,0xE))
211 RIJNDAEL_EXP_TABLE = (0x01, 0x03, 0x05, 0x0F, 0x11, 0x33, 0x55, 0xFF, 0x1A, 0x2E, 0x72, 0x96, 0xA1, 0xF8, 0x13, 0x35,
212                       0x5F, 0xE1, 0x38, 0x48, 0xD8, 0x73, 0x95, 0xA4, 0xF7, 0x02, 0x06, 0x0A, 0x1E, 0x22, 0x66, 0xAA,
213                       0xE5, 0x34, 0x5C, 0xE4, 0x37, 0x59, 0xEB, 0x26, 0x6A, 0xBE, 0xD9, 0x70, 0x90, 0xAB, 0xE6, 0x31,
214                       0x53, 0xF5, 0x04, 0x0C, 0x14, 0x3C, 0x44, 0xCC, 0x4F, 0xD1, 0x68, 0xB8, 0xD3, 0x6E, 0xB2, 0xCD,
215                       0x4C, 0xD4, 0x67, 0xA9, 0xE0, 0x3B, 0x4D, 0xD7, 0x62, 0xA6, 0xF1, 0x08, 0x18, 0x28, 0x78, 0x88,
216                       0x83, 0x9E, 0xB9, 0xD0, 0x6B, 0xBD, 0xDC, 0x7F, 0x81, 0x98, 0xB3, 0xCE, 0x49, 0xDB, 0x76, 0x9A,
217                       0xB5, 0xC4, 0x57, 0xF9, 0x10, 0x30, 0x50, 0xF0, 0x0B, 0x1D, 0x27, 0x69, 0xBB, 0xD6, 0x61, 0xA3,
218                       0xFE, 0x19, 0x2B, 0x7D, 0x87, 0x92, 0xAD, 0xEC, 0x2F, 0x71, 0x93, 0xAE, 0xE9, 0x20, 0x60, 0xA0,
219                       0xFB, 0x16, 0x3A, 0x4E, 0xD2, 0x6D, 0xB7, 0xC2, 0x5D, 0xE7, 0x32, 0x56, 0xFA, 0x15, 0x3F, 0x41,
220                       0xC3, 0x5E, 0xE2, 0x3D, 0x47, 0xC9, 0x40, 0xC0, 0x5B, 0xED, 0x2C, 0x74, 0x9C, 0xBF, 0xDA, 0x75,
221                       0x9F, 0xBA, 0xD5, 0x64, 0xAC, 0xEF, 0x2A, 0x7E, 0x82, 0x9D, 0xBC, 0xDF, 0x7A, 0x8E, 0x89, 0x80,
222                       0x9B, 0xB6, 0xC1, 0x58, 0xE8, 0x23, 0x65, 0xAF, 0xEA, 0x25, 0x6F, 0xB1, 0xC8, 0x43, 0xC5, 0x54,
223                       0xFC, 0x1F, 0x21, 0x63, 0xA5, 0xF4, 0x07, 0x09, 0x1B, 0x2D, 0x77, 0x99, 0xB0, 0xCB, 0x46, 0xCA,
224                       0x45, 0xCF, 0x4A, 0xDE, 0x79, 0x8B, 0x86, 0x91, 0xA8, 0xE3, 0x3E, 0x42, 0xC6, 0x51, 0xF3, 0x0E,
225                       0x12, 0x36, 0x5A, 0xEE, 0x29, 0x7B, 0x8D, 0x8C, 0x8F, 0x8A, 0x85, 0x94, 0xA7, 0xF2, 0x0D, 0x17,
226                       0x39, 0x4B, 0xDD, 0x7C, 0x84, 0x97, 0xA2, 0xFD, 0x1C, 0x24, 0x6C, 0xB4, 0xC7, 0x52, 0xF6, 0x01)
227 RIJNDAEL_LOG_TABLE = (0x00, 0x00, 0x19, 0x01, 0x32, 0x02, 0x1a, 0xc6, 0x4b, 0xc7, 0x1b, 0x68, 0x33, 0xee, 0xdf, 0x03,
228                       0x64, 0x04, 0xe0, 0x0e, 0x34, 0x8d, 0x81, 0xef, 0x4c, 0x71, 0x08, 0xc8, 0xf8, 0x69, 0x1c, 0xc1,
229                       0x7d, 0xc2, 0x1d, 0xb5, 0xf9, 0xb9, 0x27, 0x6a, 0x4d, 0xe4, 0xa6, 0x72, 0x9a, 0xc9, 0x09, 0x78,
230                       0x65, 0x2f, 0x8a, 0x05, 0x21, 0x0f, 0xe1, 0x24, 0x12, 0xf0, 0x82, 0x45, 0x35, 0x93, 0xda, 0x8e,
231                       0x96, 0x8f, 0xdb, 0xbd, 0x36, 0xd0, 0xce, 0x94, 0x13, 0x5c, 0xd2, 0xf1, 0x40, 0x46, 0x83, 0x38,
232                       0x66, 0xdd, 0xfd, 0x30, 0xbf, 0x06, 0x8b, 0x62, 0xb3, 0x25, 0xe2, 0x98, 0x22, 0x88, 0x91, 0x10,
233                       0x7e, 0x6e, 0x48, 0xc3, 0xa3, 0xb6, 0x1e, 0x42, 0x3a, 0x6b, 0x28, 0x54, 0xfa, 0x85, 0x3d, 0xba,
234                       0x2b, 0x79, 0x0a, 0x15, 0x9b, 0x9f, 0x5e, 0xca, 0x4e, 0xd4, 0xac, 0xe5, 0xf3, 0x73, 0xa7, 0x57,
235                       0xaf, 0x58, 0xa8, 0x50, 0xf4, 0xea, 0xd6, 0x74, 0x4f, 0xae, 0xe9, 0xd5, 0xe7, 0xe6, 0xad, 0xe8,
236                       0x2c, 0xd7, 0x75, 0x7a, 0xeb, 0x16, 0x0b, 0xf5, 0x59, 0xcb, 0x5f, 0xb0, 0x9c, 0xa9, 0x51, 0xa0,
237                       0x7f, 0x0c, 0xf6, 0x6f, 0x17, 0xc4, 0x49, 0xec, 0xd8, 0x43, 0x1f, 0x2d, 0xa4, 0x76, 0x7b, 0xb7,
238                       0xcc, 0xbb, 0x3e, 0x5a, 0xfb, 0x60, 0xb1, 0x86, 0x3b, 0x52, 0xa1, 0x6c, 0xaa, 0x55, 0x29, 0x9d,
239                       0x97, 0xb2, 0x87, 0x90, 0x61, 0xbe, 0xdc, 0xfc, 0xbc, 0x95, 0xcf, 0xcd, 0x37, 0x3f, 0x5b, 0xd1,
240                       0x53, 0x39, 0x84, 0x3c, 0x41, 0xa2, 0x6d, 0x47, 0x14, 0x2a, 0x9e, 0x5d, 0x56, 0xf2, 0xd3, 0xab,
241                       0x44, 0x11, 0x92, 0xd9, 0x23, 0x20, 0x2e, 0x89, 0xb4, 0x7c, 0xb8, 0x26, 0x77, 0x99, 0xe3, 0xa5,
242                       0x67, 0x4a, 0xed, 0xde, 0xc5, 0x31, 0xfe, 0x18, 0x0d, 0x63, 0x8c, 0x80, 0xc0, 0xf7, 0x70, 0x07)
243
244 def sub_bytes(data):
245     return [SBOX[x] for x in data]
246
247 def sub_bytes_inv(data):
248     return [SBOX_INV[x] for x in data]
249
250 def rotate(data):
251     return data[1:] + [data[0]]
252
253 def key_schedule_core(data, rcon_iteration):
254     data = rotate(data)
255     data = sub_bytes(data)
256     data[0] = data[0] ^ RCON[rcon_iteration]
257     
258     return data
259
260 def xor(data1, data2):
261     return [x^y for x, y in zip(data1, data2)]
262
263 def rijndael_mul(a, b):
264     if(a==0 or b==0):
265         return 0
266     return RIJNDAEL_EXP_TABLE[(RIJNDAEL_LOG_TABLE[a] + RIJNDAEL_LOG_TABLE[b]) % 0xFF]
267
268 def mix_column(data, matrix):
269     data_mixed = []
270     for row in range(4):
271         mixed = 0
272         for column in range(4):
273             # xor is (+) and (-)
274             mixed ^= rijndael_mul(data[column], matrix[row][column])
275         data_mixed.append(mixed)
276     return data_mixed
277
278 def mix_columns(data, matrix=MIX_COLUMN_MATRIX):
279     data_mixed = []
280     for i in range(4):
281         column = data[i*4 : (i+1)*4]
282         data_mixed += mix_column(column, matrix)
283     return data_mixed
284
285 def mix_columns_inv(data):
286     return mix_columns(data, MIX_COLUMN_MATRIX_INV)
287
288 def shift_rows(data):
289     data_shifted = []
290     for column in range(4):
291         for row in range(4):
292             data_shifted.append( data[((column + row) & 0b11) * 4 + row] )
293     return data_shifted
294
295 def shift_rows_inv(data):
296     data_shifted = []
297     for column in range(4):
298         for row in range(4):
299             data_shifted.append( data[((column - row) & 0b11) * 4 + row] )
300     return data_shifted
301
302 def inc(data):
303     data = data[:] # copy
304     for i in range(len(data)-1,-1,-1):
305         if data[i] == 255:
306             data[i] = 0
307         else:
308             data[i] = data[i] + 1
309             break
310     return data