diff --git a/extmod/modtrezorcrypto/modtrezorcrypto-aes.h b/extmod/modtrezorcrypto/modtrezorcrypto-aes.h new file mode 100644 index 0000000000..c54636258c --- /dev/null +++ b/extmod/modtrezorcrypto/modtrezorcrypto-aes.h @@ -0,0 +1,141 @@ +/* + * Copyright (c) Pavol Rusnak, SatoshiLabs + * + * Licensed under TREZOR License + * see LICENSE file for details + */ + +#include "py/objstr.h" + +#include "trezor-crypto/aes.h" + +/* +*/ + +typedef struct _mp_obj_AES_t { + mp_obj_base_t base; + union { + aes_encrypt_ctx encrypt_ctx; + aes_decrypt_ctx decrypt_ctx; + } ctx; + mp_int_t mode; + uint8_t iv[AES_BLOCK_SIZE]; + uint8_t ctr[AES_BLOCK_SIZE]; +} mp_obj_AES_t; + +/// def trezor.crypto.aes.AES(mode:int, key: bytes, iv: bytes=None) -> AES +STATIC mp_obj_t mod_TrezorCrypto_AES_make_new(const mp_obj_type_t *type, size_t n_args, size_t n_kw, const mp_obj_t *args) { + mp_arg_check_num(n_args, n_kw, 2, 3, false); + mp_obj_AES_t *o = m_new_obj(mp_obj_AES_t); + o->base.type = type; + o->mode = mp_obj_get_int(args[0]); + if ((o->mode & 0x7F) > 0x04) { + nlr_raise(mp_obj_new_exception_msg(&mp_type_ValueError, "Invalid AES mode")); + } + mp_buffer_info_t key; + mp_get_buffer_raise(args[1], &key, MP_BUFFER_READ); + if (key.len != 16 && key.len != 24 && key.len != 32) { + nlr_raise(mp_obj_new_exception_msg(&mp_type_ValueError, "Invalid length of key (has to be 128, 192 or 256 bits)")); + } + if (n_args > 2) { + mp_buffer_info_t iv; + mp_get_buffer_raise(args[2], &iv, MP_BUFFER_READ); + if (iv.len != AES_BLOCK_SIZE) { + nlr_raise(mp_obj_new_exception_msg(&mp_type_ValueError, "Invalid length of initialization vector (has to be 128 bits)")); + } + memcpy(o->iv, iv.buf, AES_BLOCK_SIZE); + } else { + memset(o->iv, 0, AES_BLOCK_SIZE); + } + memset(o->ctr, 0, AES_BLOCK_SIZE); + switch (key.len) { + case 16: + if ((o->mode & 0x80) == 0x00) { + aes_encrypt_key128(key.buf, &(o->ctx.encrypt_ctx)); + } else { + aes_decrypt_key128(key.buf, &(o->ctx.decrypt_ctx)); + } + break; + case 24: + if ((o->mode & 0x80) == 0x00) { + aes_encrypt_key192(key.buf, &(o->ctx.encrypt_ctx)); + } else { + aes_decrypt_key192(key.buf, &(o->ctx.decrypt_ctx)); + } + break; + case 32: + if ((o->mode & 0x80) == 0x00) { + aes_encrypt_key256(key.buf, &(o->ctx.encrypt_ctx)); + } else { + aes_decrypt_key256(key.buf, &(o->ctx.decrypt_ctx)); + } + break; + } + return MP_OBJ_FROM_PTR(o); +} + +/// def AES.crypt(self, data: bytes) -> bytes +STATIC mp_obj_t mod_TrezorCrypto_AES_update(mp_obj_t self, mp_obj_t data) { + mp_buffer_info_t buf; + mp_get_buffer_raise(data, &buf, MP_BUFFER_READ); + mp_obj_AES_t *o = MP_OBJ_TO_PTR(self); + vstr_t vstr; + vstr_init_len(&vstr, buf.len); + switch (o->mode & 0x7F) { + case 0x00: // ECB + if (buf.len & (AES_BLOCK_SIZE - 1)) { + nlr_raise(mp_obj_new_exception_msg(&mp_type_ValueError, "Invalid data length")); + } + if ((o->mode & 0x80) == 0x00) { + aes_ecb_encrypt(buf.buf, (unsigned char *)vstr.buf, buf.len, &(o->ctx.encrypt_ctx)); + } else { + aes_ecb_decrypt(buf.buf, (unsigned char *)vstr.buf, buf.len, &(o->ctx.decrypt_ctx)); + } + break; + case 0x01: // CBC + if (buf.len & (AES_BLOCK_SIZE - 1)) { + nlr_raise(mp_obj_new_exception_msg(&mp_type_ValueError, "Invalid data length")); + } + if ((o->mode & 0x80) == 0x00) { + aes_cbc_encrypt(buf.buf, (unsigned char *)vstr.buf, buf.len, o->iv, &(o->ctx.encrypt_ctx)); + } else { + aes_cbc_decrypt(buf.buf, (unsigned char *)vstr.buf, buf.len, o->iv, &(o->ctx.decrypt_ctx)); + } + break; + case 0x02: // CFB + if ((o->mode & 0x80) == 0x00) { + aes_cfb_encrypt(buf.buf, (unsigned char *)vstr.buf, buf.len, o->iv, &(o->ctx.encrypt_ctx)); + } else { + aes_cfb_decrypt(buf.buf, (unsigned char *)vstr.buf, buf.len, o->iv, &(o->ctx.encrypt_ctx)); + } + break; + case 0x03: // OFB (encrypt == decrypt) + aes_ofb_crypt(buf.buf, (unsigned char *)vstr.buf, buf.len, o->iv, &(o->ctx.encrypt_ctx)); + break; + case 0x04: // CTR (encrypt == decrypt) + aes_ctr_crypt(buf.buf, (unsigned char *)vstr.buf, buf.len, o->ctr, aes_ctr_cbuf_inc, &(o->ctx.encrypt_ctx)); + break; + } + return mp_obj_new_str_from_vstr(&mp_type_bytes, &vstr); +} +STATIC MP_DEFINE_CONST_FUN_OBJ_2(mod_TrezorCrypto_AES_update_obj, mod_TrezorCrypto_AES_update); + +STATIC mp_obj_t mod_TrezorCrypto_AES___del__(mp_obj_t self) { + mp_obj_AES_t *o = MP_OBJ_TO_PTR(self); + memset(&(o->ctx), 0, sizeof(aes_encrypt_ctx)); + return mp_const_none; +} +STATIC MP_DEFINE_CONST_FUN_OBJ_1(mod_TrezorCrypto_AES___del___obj, mod_TrezorCrypto_AES___del__); + +STATIC const mp_rom_map_elem_t mod_TrezorCrypto_AES_locals_dict_table[] = { + { MP_ROM_QSTR(MP_QSTR_update), MP_ROM_PTR(&mod_TrezorCrypto_AES_update_obj) }, + { MP_ROM_QSTR(MP_QSTR___del__), MP_ROM_PTR(&mod_TrezorCrypto_AES___del___obj) }, +}; +STATIC MP_DEFINE_CONST_DICT(mod_TrezorCrypto_AES_locals_dict, mod_TrezorCrypto_AES_locals_dict_table); + +STATIC const mp_obj_type_t mod_TrezorCrypto_AES_type = { + { &mp_type_type }, + .name = MP_QSTR_AES, + .make_new = mod_TrezorCrypto_AES_make_new, + .locals_dict = (void*)&mod_TrezorCrypto_AES_locals_dict, +}; diff --git a/extmod/modtrezorcrypto/modtrezorcrypto.c b/extmod/modtrezorcrypto/modtrezorcrypto.c index 02a047fb2d..0125b4997c 100644 --- a/extmod/modtrezorcrypto/modtrezorcrypto.c +++ b/extmod/modtrezorcrypto/modtrezorcrypto.c @@ -15,6 +15,7 @@ #if MICROPY_PY_TREZORCRYPTO +#include "modtrezorcrypto-aes.h" #include "modtrezorcrypto-bip39.h" #include "modtrezorcrypto-ed25519.h" #include "modtrezorcrypto-pbkdf2.h" @@ -30,6 +31,7 @@ STATIC const mp_rom_map_elem_t mp_module_TrezorCrypto_globals_table[] = { { MP_ROM_QSTR(MP_QSTR___name__), MP_ROM_QSTR(MP_QSTR_TrezorCrypto) }, + { MP_ROM_QSTR(MP_QSTR_AES), MP_ROM_PTR(&mod_TrezorCrypto_AES_type) }, { MP_ROM_QSTR(MP_QSTR_Bip39), MP_ROM_PTR(&mod_TrezorCrypto_Bip39_type) }, { MP_ROM_QSTR(MP_QSTR_Ed25519), MP_ROM_PTR(&mod_TrezorCrypto_Ed25519_type) }, { MP_ROM_QSTR(MP_QSTR_Nist256p1), MP_ROM_PTR(&mod_TrezorCrypto_Nist256p1_type) }, diff --git a/src/trezor/crypto/aes.py b/src/trezor/crypto/aes.py new file mode 100644 index 0000000000..24c05f7d51 --- /dev/null +++ b/src/trezor/crypto/aes.py @@ -0,0 +1,31 @@ +from TrezorCrypto import AES as _AES + +def AES_ECB_Encrypt(key): + return _AES(0x00, key) + +def AES_ECB_Decrypt(key): + return _AES(0x80, key) + +def AES_CBC_Encrypt(key, iv): + return _AES(0x01, key, iv) + +def AES_CBC_Decrypt(key, iv): + return _AES(0x81, key, iv) + +def AES_CFB_Encrypt(key, iv): + return _AES(0x02, key, iv) + +def AES_CFB_Decrypt(key, iv): + return _AES(0x82, key, iv) + +def AES_OFB_Encrypt(key, iv): + return _AES(0x03, key, iv) + +def AES_OFB_Decrypt(key, iv): + return _AES(0x83, key, iv) + +def AES_CTR_Encrypt(key): + return _AES(0x04, key) + +def AES_CTR_Decrypt(key): + return _AES(0x84, key) diff --git a/vendor/micropython b/vendor/micropython index 3f91c5658d..25e1380cc1 160000 --- a/vendor/micropython +++ b/vendor/micropython @@ -1 +1 @@ -Subproject commit 3f91c5658d20547d7d5211f2010b3aeabd32288e +Subproject commit 25e1380cc1fc5197de8f05d6b8df02a30e8d3f65