1/*
2 * This file is part of the MicroPython project, http://micropython.org/
3 *
4 * The MIT License (MIT)
5 *
6 * Copyright (c) 2017-2018 Paul Sokolovsky
7 * Copyright (c) 2018 Yonatan Goldschmidt
8 *
9 * Permission is hereby granted, free of charge, to any person obtaining a copy
10 * of this software and associated documentation files (the "Software"), to deal
11 * in the Software without restriction, including without limitation the rights
12 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
13 * copies of the Software, and to permit persons to whom the Software is
14 * furnished to do so, subject to the following conditions:
15 *
16 * The above copyright notice and this permission notice shall be included in
17 * all copies or substantial portions of the Software.
18 *
19 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
20 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
21 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
22 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
23 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
24 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
25 * THE SOFTWARE.
26 */
27
28#include "py/mpconfig.h"
29
30#if MICROPY_PY_UCRYPTOLIB
31
32#include <assert.h>
33#include <string.h>
34
35#include "py/runtime.h"
36
37// This module implements crypto ciphers API, roughly following
38// https://www.python.org/dev/peps/pep-0272/ . Exact implementation
39// of PEP 272 can be made with a simple wrapper which adds all the
40// needed boilerplate.
41
42// values follow PEP 272
43enum {
44 UCRYPTOLIB_MODE_ECB = 1,
45 UCRYPTOLIB_MODE_CBC = 2,
46 UCRYPTOLIB_MODE_CTR = 6,
47};
48
49struct ctr_params {
50 // counter is the IV of the AES context.
51
52 size_t offset; // in encrypted_counter
53 // encrypted counter
54 uint8_t encrypted_counter[16];
55};
56
57#if MICROPY_SSL_AXTLS
58#include "lib/axtls/crypto/crypto.h"
59
60#define AES_CTX_IMPL AES_CTX
61#endif
62
63#if MICROPY_SSL_MBEDTLS
64#include <mbedtls/aes.h>
65
66// we can't run mbedtls AES key schedule until we know whether we're used for encrypt or decrypt.
67// therefore, we store the key & keysize and on the first call to encrypt/decrypt we override them
68// with the mbedtls_aes_context, as they are not longer required. (this is done to save space)
69struct mbedtls_aes_ctx_with_key {
70 union {
71 mbedtls_aes_context mbedtls_ctx;
72 struct {
73 uint8_t key[32];
74 uint8_t keysize;
75 } init_data;
76 } u;
77 unsigned char iv[16];
78};
79#define AES_CTX_IMPL struct mbedtls_aes_ctx_with_key
80#endif
81
82typedef struct _mp_obj_aes_t {
83 mp_obj_base_t base;
84 AES_CTX_IMPL ctx;
85 uint8_t block_mode : 6;
86#define AES_KEYTYPE_NONE 0
87#define AES_KEYTYPE_ENC 1
88#define AES_KEYTYPE_DEC 2
89 uint8_t key_type : 2;
90} mp_obj_aes_t;
91
92static inline bool is_ctr_mode(int block_mode) {
93 #if MICROPY_PY_UCRYPTOLIB_CTR
94 return block_mode == UCRYPTOLIB_MODE_CTR;
95 #else
96 return false;
97 #endif
98}
99
100static inline struct ctr_params *ctr_params_from_aes(mp_obj_aes_t *o) {
101 // ctr_params follows aes object struct
102 return (struct ctr_params *)&o[1];
103}
104
105#if MICROPY_SSL_AXTLS
106STATIC void aes_initial_set_key_impl(AES_CTX_IMPL *ctx, const uint8_t *key, size_t keysize, const uint8_t iv[16]) {
107 assert(16 == keysize || 32 == keysize);
108 AES_set_key(ctx, key, iv, (16 == keysize) ? AES_MODE_128 : AES_MODE_256);
109}
110
111STATIC void aes_final_set_key_impl(AES_CTX_IMPL *ctx, bool encrypt) {
112 if (!encrypt) {
113 AES_convert_key(ctx);
114 }
115}
116
117STATIC void aes_process_ecb_impl(AES_CTX_IMPL *ctx, const uint8_t in[16], uint8_t out[16], bool encrypt) {
118 memcpy(out, in, 16);
119 // We assume that out (vstr.buf or given output buffer) is uint32_t aligned
120 uint32_t *p = (uint32_t *)out;
121 // axTLS likes it weird and complicated with byteswaps
122 for (int i = 0; i < 4; i++) {
123 p[i] = MP_HTOBE32(p[i]);
124 }
125 if (encrypt) {
126 AES_encrypt(ctx, p);
127 } else {
128 AES_decrypt(ctx, p);
129 }
130 for (int i = 0; i < 4; i++) {
131 p[i] = MP_BE32TOH(p[i]);
132 }
133}
134
135STATIC void aes_process_cbc_impl(AES_CTX_IMPL *ctx, const uint8_t *in, uint8_t *out, size_t in_len, bool encrypt) {
136 if (encrypt) {
137 AES_cbc_encrypt(ctx, in, out, in_len);
138 } else {
139 AES_cbc_decrypt(ctx, in, out, in_len);
140 }
141}
142
143#if MICROPY_PY_UCRYPTOLIB_CTR
144// axTLS doesn't have CTR support out of the box. This implements the counter part using the ECB primitive.
145STATIC void aes_process_ctr_impl(AES_CTX_IMPL *ctx, const uint8_t *in, uint8_t *out, size_t in_len, struct ctr_params *ctr_params) {
146 size_t n = ctr_params->offset;
147 uint8_t *const counter = ctx->iv;
148
149 while (in_len--) {
150 if (n == 0) {
151 aes_process_ecb_impl(ctx, counter, ctr_params->encrypted_counter, true);
152
153 // increment the 128-bit counter
154 for (int i = 15; i >= 0; --i) {
155 if (++counter[i] != 0) {
156 break;
157 }
158 }
159 }
160
161 *out++ = *in++ ^ ctr_params->encrypted_counter[n];
162 n = (n + 1) & 0xf;
163 }
164
165 ctr_params->offset = n;
166}
167#endif
168
169#endif
170
171#if MICROPY_SSL_MBEDTLS
172STATIC void aes_initial_set_key_impl(AES_CTX_IMPL *ctx, const uint8_t *key, size_t keysize, const uint8_t iv[16]) {
173 ctx->u.init_data.keysize = keysize;
174 memcpy(ctx->u.init_data.key, key, keysize);
175
176 if (NULL != iv) {
177 memcpy(ctx->iv, iv, sizeof(ctx->iv));
178 }
179}
180
181STATIC void aes_final_set_key_impl(AES_CTX_IMPL *ctx, bool encrypt) {
182 // first, copy key aside
183 uint8_t key[32];
184 uint8_t keysize = ctx->u.init_data.keysize;
185 memcpy(key, ctx->u.init_data.key, keysize);
186 // now, override key with the mbedtls context object
187 mbedtls_aes_init(&ctx->u.mbedtls_ctx);
188
189 // setkey call will succeed, we've already checked the keysize earlier.
190 assert(16 == keysize || 32 == keysize);
191 if (encrypt) {
192 mbedtls_aes_setkey_enc(&ctx->u.mbedtls_ctx, key, keysize * 8);
193 } else {
194 mbedtls_aes_setkey_dec(&ctx->u.mbedtls_ctx, key, keysize * 8);
195 }
196}
197
198STATIC void aes_process_ecb_impl(AES_CTX_IMPL *ctx, const uint8_t in[16], uint8_t out[16], bool encrypt) {
199 mbedtls_aes_crypt_ecb(&ctx->u.mbedtls_ctx, encrypt ? MBEDTLS_AES_ENCRYPT : MBEDTLS_AES_DECRYPT, in, out);
200}
201
202STATIC void aes_process_cbc_impl(AES_CTX_IMPL *ctx, const uint8_t *in, uint8_t *out, size_t in_len, bool encrypt) {
203 mbedtls_aes_crypt_cbc(&ctx->u.mbedtls_ctx, encrypt ? MBEDTLS_AES_ENCRYPT : MBEDTLS_AES_DECRYPT, in_len, ctx->iv, in, out);
204}
205
206#if MICROPY_PY_UCRYPTOLIB_CTR
207STATIC void aes_process_ctr_impl(AES_CTX_IMPL *ctx, const uint8_t *in, uint8_t *out, size_t in_len, struct ctr_params *ctr_params) {
208 mbedtls_aes_crypt_ctr(&ctx->u.mbedtls_ctx, in_len, &ctr_params->offset, ctx->iv, ctr_params->encrypted_counter, in, out);
209}
210#endif
211
212#endif
213
214STATIC mp_obj_t ucryptolib_aes_make_new(const mp_obj_type_t *type, size_t n_args, size_t n_kw, const mp_obj_t *args) {
215 mp_arg_check_num(n_args, n_kw, 2, 3, false);
216
217 const mp_int_t block_mode = mp_obj_get_int(args[1]);
218
219 switch (block_mode) {
220 case UCRYPTOLIB_MODE_ECB:
221 case UCRYPTOLIB_MODE_CBC:
222 #if MICROPY_PY_UCRYPTOLIB_CTR
223 case UCRYPTOLIB_MODE_CTR:
224 #endif
225 break;
226
227 default:
228 mp_raise_ValueError(MP_ERROR_TEXT("mode"));
229 }
230
231 mp_obj_aes_t *o = m_new_obj_var(mp_obj_aes_t, struct ctr_params, !!is_ctr_mode(block_mode));
232 o->base.type = type;
233
234 o->block_mode = block_mode;
235 o->key_type = AES_KEYTYPE_NONE;
236
237 mp_buffer_info_t keyinfo;
238 mp_get_buffer_raise(args[0], &keyinfo, MP_BUFFER_READ);
239 if (32 != keyinfo.len && 16 != keyinfo.len) {
240 mp_raise_ValueError(MP_ERROR_TEXT("key"));
241 }
242
243 mp_buffer_info_t ivinfo;
244 ivinfo.buf = NULL;
245 if (n_args > 2 && args[2] != mp_const_none) {
246 mp_get_buffer_raise(args[2], &ivinfo, MP_BUFFER_READ);
247
248 if (16 != ivinfo.len) {
249 mp_raise_ValueError(MP_ERROR_TEXT("IV"));
250 }
251 } else if (o->block_mode == UCRYPTOLIB_MODE_CBC || is_ctr_mode(o->block_mode)) {
252 mp_raise_ValueError(MP_ERROR_TEXT("IV"));
253 }
254
255 if (is_ctr_mode(block_mode)) {
256 ctr_params_from_aes(o)->offset = 0;
257 }
258
259 aes_initial_set_key_impl(&o->ctx, keyinfo.buf, keyinfo.len, ivinfo.buf);
260
261 return MP_OBJ_FROM_PTR(o);
262}
263
264STATIC mp_obj_t aes_process(size_t n_args, const mp_obj_t *args, bool encrypt) {
265 mp_obj_aes_t *self = MP_OBJ_TO_PTR(args[0]);
266
267 mp_obj_t in_buf = args[1];
268 mp_obj_t out_buf = MP_OBJ_NULL;
269 if (n_args > 2) {
270 out_buf = args[2];
271 }
272
273 mp_buffer_info_t in_bufinfo;
274 mp_get_buffer_raise(in_buf, &in_bufinfo, MP_BUFFER_READ);
275
276 if (!is_ctr_mode(self->block_mode) && in_bufinfo.len % 16 != 0) {
277 mp_raise_ValueError(MP_ERROR_TEXT("blksize % 16"));
278 }
279
280 vstr_t vstr;
281 mp_buffer_info_t out_bufinfo;
282 uint8_t *out_buf_ptr;
283
284 if (out_buf != MP_OBJ_NULL) {
285 mp_get_buffer_raise(out_buf, &out_bufinfo, MP_BUFFER_WRITE);
286 if (out_bufinfo.len < in_bufinfo.len) {
287 mp_raise_ValueError(MP_ERROR_TEXT("output too small"));
288 }
289 out_buf_ptr = out_bufinfo.buf;
290 } else {
291 vstr_init_len(&vstr, in_bufinfo.len);
292 out_buf_ptr = (uint8_t *)vstr.buf;
293 }
294
295 if (AES_KEYTYPE_NONE == self->key_type) {
296 // always set key for encryption if CTR mode.
297 const bool encrypt_mode = encrypt || is_ctr_mode(self->block_mode);
298 aes_final_set_key_impl(&self->ctx, encrypt_mode);
299 self->key_type = encrypt ? AES_KEYTYPE_ENC : AES_KEYTYPE_DEC;
300 } else {
301 if ((encrypt && self->key_type == AES_KEYTYPE_DEC) ||
302 (!encrypt && self->key_type == AES_KEYTYPE_ENC)) {
303
304 mp_raise_ValueError(MP_ERROR_TEXT("can't encrypt & decrypt"));
305 }
306 }
307
308 switch (self->block_mode) {
309 case UCRYPTOLIB_MODE_ECB: {
310 uint8_t *in = in_bufinfo.buf, *out = out_buf_ptr;
311 uint8_t *top = in + in_bufinfo.len;
312 for (; in < top; in += 16, out += 16) {
313 aes_process_ecb_impl(&self->ctx, in, out, encrypt);
314 }
315 break;
316 }
317
318 case UCRYPTOLIB_MODE_CBC:
319 aes_process_cbc_impl(&self->ctx, in_bufinfo.buf, out_buf_ptr, in_bufinfo.len, encrypt);
320 break;
321
322 #if MICROPY_PY_UCRYPTOLIB_CTR
323 case UCRYPTOLIB_MODE_CTR:
324 aes_process_ctr_impl(&self->ctx, in_bufinfo.buf, out_buf_ptr, in_bufinfo.len,
325 ctr_params_from_aes(self));
326 break;
327 #endif
328 }
329
330 if (out_buf != MP_OBJ_NULL) {
331 return out_buf;
332 }
333 return mp_obj_new_str_from_vstr(&mp_type_bytes, &vstr);
334}
335
336STATIC mp_obj_t ucryptolib_aes_encrypt(size_t n_args, const mp_obj_t *args) {
337 return aes_process(n_args, args, true);
338}
339STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(ucryptolib_aes_encrypt_obj, 2, 3, ucryptolib_aes_encrypt);
340
341STATIC mp_obj_t ucryptolib_aes_decrypt(size_t n_args, const mp_obj_t *args) {
342 return aes_process(n_args, args, false);
343}
344STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(ucryptolib_aes_decrypt_obj, 2, 3, ucryptolib_aes_decrypt);
345
346STATIC const mp_rom_map_elem_t ucryptolib_aes_locals_dict_table[] = {
347 { MP_ROM_QSTR(MP_QSTR_encrypt), MP_ROM_PTR(&ucryptolib_aes_encrypt_obj) },
348 { MP_ROM_QSTR(MP_QSTR_decrypt), MP_ROM_PTR(&ucryptolib_aes_decrypt_obj) },
349};
350STATIC MP_DEFINE_CONST_DICT(ucryptolib_aes_locals_dict, ucryptolib_aes_locals_dict_table);
351
352STATIC const mp_obj_type_t ucryptolib_aes_type = {
353 { &mp_type_type },
354 .name = MP_QSTR_aes,
355 .make_new = ucryptolib_aes_make_new,
356 .locals_dict = (void *)&ucryptolib_aes_locals_dict,
357};
358
359STATIC const mp_rom_map_elem_t mp_module_ucryptolib_globals_table[] = {
360 { MP_ROM_QSTR(MP_QSTR___name__), MP_ROM_QSTR(MP_QSTR_ucryptolib) },
361 { MP_ROM_QSTR(MP_QSTR_aes), MP_ROM_PTR(&ucryptolib_aes_type) },
362 #if MICROPY_PY_UCRYPTOLIB_CONSTS
363 { MP_ROM_QSTR(MP_QSTR_MODE_ECB), MP_ROM_INT(UCRYPTOLIB_MODE_ECB) },
364 { MP_ROM_QSTR(MP_QSTR_MODE_CBC), MP_ROM_INT(UCRYPTOLIB_MODE_CBC) },
365 #if MICROPY_PY_UCRYPTOLIB_CTR
366 { MP_ROM_QSTR(MP_QSTR_MODE_CTR), MP_ROM_INT(UCRYPTOLIB_MODE_CTR) },
367 #endif
368 #endif
369};
370
371STATIC MP_DEFINE_CONST_DICT(mp_module_ucryptolib_globals, mp_module_ucryptolib_globals_table);
372
373const mp_obj_module_t mp_module_ucryptolib = {
374 .base = { &mp_type_module },
375 .globals = (mp_obj_dict_t *)&mp_module_ucryptolib_globals,
376};
377
378#endif // MICROPY_PY_UCRYPTOLIB
379