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