aesio: add basic AES encryption and decryption

This adds initial support for an AES module named aesio.  This
implementation supports only a subset of AES modes, namely
ECB, CBC, and CTR modes.

Example usage:

```
>>> import aesio
>>>
>>> key = b'Sixteen byte key'
>>> cipher = aesio.AES(key, aesio.MODE_ECB)
>>> output = bytearray(16)
>>> cipher.encrypt_into(b'Circuit Python!!', output)
>>> output
bytearray(b'E\x14\x85\x18\x9a\x9c\r\x95>\xa7kV\xa2`\x8b\n')
>>>
```

This key is 16-bytes, so it uses AES128.  If your key is 24- or 32-
bytes long, it will switch to AES192 or AES256 respectively.

This has been tested with many of the official NIST test vectors,
such as those used in `pycryptodome` at
https://github.com/Legrandin/pycryptodome/tree/39626a5b01ce5c1cf51d022be166ad0aea722177/lib/Crypto/SelfTest/Cipher/test_vectors/AES

CTR has not been tested as NIST does not provide test vectors for it.

Signed-off-by: Sean Cross <sean@xobs.io>

3 files changed, 403 insertions, 0 deletions

diff --git a/shared-bindings/aesio/__init__.c b/shared-bindings/aesio/__init__.c
new file mode 100644
index 000000000..43dc73e8b
--- /dev/null
+++ b/shared-bindings/aesio/__init__.c
@@ -0,0 +1,79 @@
+/*
+ * This file is part of the MicroPython project, http://micropython.org/
+ *
+ * The MIT License (MIT)
+ *
+ * Copyright (c) 2017 Scott Shawcroft for Adafruit Industries
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ * THE SOFTWARE.
+ */
+
+#include <stdint.h>
+
+#include "py/obj.h"
+#include "py/runtime.h"
+
+#include "__init__.h"
+
+//| :mod:`aesio` --- AES encryption routines
+//| ========================================
+//|
+//| .. module:: aesio
+//|   :synopsis: Embedded implementation of AES
+//|
+//| The `AES` module contains classes used to implement encryption
+//| and decryption. It aims to be low overhead in terms of memory.
+//|
+//|
+
+//| Libraries
+//|
+//| .. toctree::
+//|     :maxdepth: 3
+//|
+//|     aes
+
+
+STATIC const mp_obj_tuple_t mp_aes_key_size_obj = {
+    {&mp_type_tuple},
+    3,
+    {
+        MP_OBJ_NEW_SMALL_INT(16),
+        MP_OBJ_NEW_SMALL_INT(24),
+        MP_OBJ_NEW_SMALL_INT(32),
+    }
+};
+
+STATIC const mp_rom_map_elem_t aesio_module_globals_table[] = {
+    {MP_ROM_QSTR(MP_QSTR___name__), MP_ROM_QSTR(MP_QSTR_aesio)},
+    {MP_ROM_QSTR(MP_QSTR_AES), MP_ROM_PTR(&aesio_aes_type) },
+    {MP_ROM_QSTR(MP_QSTR_MODE_ECB), MP_ROM_INT(AES_MODE_ECB)},
+    {MP_ROM_QSTR(MP_QSTR_MODE_CBC), MP_ROM_INT(AES_MODE_CBC)},
+    {MP_ROM_QSTR(MP_QSTR_MODE_CTR), MP_ROM_INT(AES_MODE_CTR)},
+    {MP_ROM_QSTR(MP_QSTR_block_size), MP_ROM_INT(AES_BLOCKLEN)},
+    {MP_ROM_QSTR(MP_QSTR_key_size), (mp_obj_t)&mp_aes_key_size_obj},
+};
+
+STATIC MP_DEFINE_CONST_DICT(aesio_module_globals, aesio_module_globals_table);
+
+const mp_obj_module_t aesio_module = {
+    .base = {&mp_type_module},
+    .globals = (mp_obj_dict_t *)&aesio_module_globals,
+};
+
diff --git a/shared-bindings/aesio/__init__.h b/shared-bindings/aesio/__init__.h
new file mode 100644
index 000000000..634fed2e5
--- /dev/null
+++ b/shared-bindings/aesio/__init__.h
@@ -0,0 +1,53 @@
+/*
+ * This file is part of the Micro Python project, http://micropython.org/
+ *
+ * The MIT License (MIT)
+ *
+ * Copyright (c) 2017 Scott Shawcroft for Adafruit Industries
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ * THE SOFTWARE.
+ */
+
+#ifndef MICROPY_INCLUDED_SHARED_BINDINGS_AESIO_H
+#define MICROPY_INCLUDED_SHARED_BINDINGS_AESIO_H
+
+#include "shared-module/aesio/__init__.h"
+
+extern const mp_obj_type_t aesio_aes_type;
+
+void common_hal_aesio_aes_construct(aesio_aes_obj_t* self,
+                              const uint8_t* key,
+                              uint32_t key_length,
+                              const uint8_t* iv,
+                              int mode,
+                              int counter);
+void common_hal_aesio_aes_rekey(aesio_aes_obj_t* self,
+                          const uint8_t* key,
+                          uint32_t key_length,
+                          const uint8_t* iv);
+void common_hal_aesio_aes_set_mode(aesio_aes_obj_t* self,
+                             int mode);
+void common_hal_aesio_aes_encrypt(aesio_aes_obj_t* self,
+                            uint8_t* buffer,
+                            size_t len);
+void common_hal_aesio_aes_decrypt(aesio_aes_obj_t* self,
+                            uint8_t* buffer,
+                            size_t len);
+
+#endif // MICROPY_INCLUDED_SHARED_BINDINGS_AESIO_H
diff --git a/shared-bindings/aesio/aes.c b/shared-bindings/aesio/aes.c
new file mode 100644
index 000000000..4ddcfa898
--- /dev/null
+++ b/shared-bindings/aesio/aes.c
@@ -0,0 +1,271 @@
+#include <stdint.h>
+#include <string.h>
+
+#include "py/obj.h"
+#include "py/objproperty.h"
+#include "py/runtime.h"
+
+#include "shared-bindings/aesio/__init__.h"
+
+// Defined at the end of this file
+
+//| .. currentmodule:: aesio
+//|
+//| :class:`aesio` -- Encrypt and decrypt AES streams
+//| =====================================================
+//|
+//| An object that represents an AES stream, including the current state.
+//|
+//| .. class:: AES(key, mode=0, iv=None, segment_size=8)
+//|
+//|   Create a new AES state with the given key.
+//|
+//|   :param bytearray key: A 16-, 24-, or 32-byte key
+//|   :param int mode: AES mode to use.  One of: AES.MODE_ECB, AES.MODE_CBC, or
+//|                    AES.MODE_CTR
+//|   :param bytearray iv: Initialization vector to use for CBC or CTR mode
+//|
+//|   Additional arguments are supported for legacy reasons.
+//|
+//|   Encrypting a string::
+//|
+//|     import aesio
+//|     from binascii import hexlify
+//|
+//|     key = b'Sixteen byte key'
+//|     inp = b'Circuit Python!!' # Note: 16-bytes long
+//|     outp = bytearray(len(inp))
+//|     cipher = aesio.AES(key, aesio.mode.MODE_ECB)
+//|     cipher.encrypt_into(inp, outp)
+//|     hexlify(outp)
+//|
+
+STATIC mp_obj_t aesio_aes_make_new(const mp_obj_type_t *type, size_t n_args,
+                                   const mp_obj_t *pos_args,
+                                   mp_map_t *kw_args) {
+  (void)type;
+  enum { ARG_key, ARG_mode, ARG_IV, ARG_counter, ARG_segment_size };
+  static const mp_arg_t allowed_args[] = {
+      {MP_QSTR_key, MP_ARG_OBJ | MP_ARG_REQUIRED},
+      {MP_QSTR_mode, MP_ARG_INT, {.u_int = AES_MODE_ECB}},
+      {MP_QSTR_IV, MP_ARG_OBJ},
+      {MP_QSTR_counter, MP_ARG_OBJ},
+      {MP_QSTR_segment_size, MP_ARG_INT, {.u_int = 8}},
+  };
+  mp_arg_val_t args[MP_ARRAY_SIZE(allowed_args)];
+
+  mp_arg_parse_all(n_args, pos_args, kw_args, MP_ARRAY_SIZE(allowed_args),
+                   allowed_args, args);
+
+  aesio_aes_obj_t *self = m_new_obj(aesio_aes_obj_t);
+  self->base.type = &aesio_aes_type;
+
+  mp_buffer_info_t bufinfo;
+
+  const uint8_t *key = NULL;
+  uint32_t key_length = 0;
+  if (mp_get_buffer(args[ARG_key].u_obj, &bufinfo, MP_BUFFER_READ)) {
+    if ((bufinfo.len != 16) && (bufinfo.len != 24) && (bufinfo.len != 32)) {
+      mp_raise_TypeError(translate("Key must be 16, 24, or 32 bytes long"));
+    }
+    key = bufinfo.buf;
+    key_length = bufinfo.len;
+  } else {
+    mp_raise_TypeError(translate("No key was specified"));
+  }
+
+  int mode = args[ARG_mode].u_int;
+  switch (args[ARG_mode].u_int) {
+  case AES_MODE_CBC:
+  case AES_MODE_ECB:
+  case AES_MODE_CTR:
+    break;
+  default:
+    mp_raise_TypeError(translate("Requested AES mode is unsupported"));
+  }
+
+  // IV is required for CBC mode and is ignored for other modes.
+  const uint8_t *iv = NULL;
+  if (args[ARG_IV].u_obj != NULL &&
+      mp_get_buffer(args[ARG_IV].u_obj, &bufinfo, MP_BUFFER_READ)) {
+    if (bufinfo.len != AES_BLOCKLEN) {
+      mp_raise_TypeError_varg(translate("IV must be %d bytes long"),
+                              AES_BLOCKLEN);
+    }
+    iv = bufinfo.buf;
+  }
+
+  common_hal_aesio_aes_construct(self, key, key_length, iv, mode,
+                                 args[ARG_counter].u_int);
+  return MP_OBJ_FROM_PTR(self);
+}
+
+STATIC mp_obj_t aesio_aes_rekey(size_t n_args, const mp_obj_t *pos_args) {
+  aesio_aes_obj_t *self = MP_OBJ_TO_PTR(pos_args[0]);
+
+  size_t key_length = 0;
+  const uint8_t *key =
+      (const uint8_t *)mp_obj_str_get_data(pos_args[1], &key_length);
+  if (key == NULL) {
+    mp_raise_ValueError(translate("No key was specified"));
+  }
+  if ((key_length != 16) && (key_length != 24) && (key_length != 32)) {
+    mp_raise_TypeError(translate("Key must be 16, 24, or 32 bytes long"));
+  }
+
+  const uint8_t *iv = NULL;
+  if (n_args > 2) {
+    size_t iv_length = 0;
+    iv = (const uint8_t *)mp_obj_str_get_data(pos_args[2], &iv_length);
+    if (iv_length != AES_BLOCKLEN) {
+      mp_raise_TypeError_varg(translate("IV must be %d bytes long"),
+                              AES_BLOCKLEN);
+    }
+  }
+
+  common_hal_aesio_aes_rekey(self, key, key_length, iv);
+  return mp_const_none;
+}
+
+MP_DEFINE_CONST_FUN_OBJ_VAR(aesio_aes_rekey_obj, 2, aesio_aes_rekey);
+
+STATIC void validate_length(aesio_aes_obj_t *self, size_t src_length,
+                            size_t dest_length) {
+  if (src_length != dest_length) {
+    mp_raise_ValueError(
+        translate("Source and destination buffers must be the same length"));
+  }
+
+  switch (self->mode) {
+  case AES_MODE_ECB:
+    if (src_length != 16) {
+      mp_raise_msg(&mp_type_ValueError,
+                   translate("ECB only operates on 16 bytes at a time"));
+    }
+    break;
+  case AES_MODE_CBC:
+    if ((src_length & 15) != 0) {
+      mp_raise_msg(&mp_type_ValueError,
+                   translate("CBC blocks must be multiples of 16 bytes"));
+    }
+    break;
+  case AES_MODE_CTR:
+    break;
+  }
+}
+
+//|   .. method:: encrypt_into(src, dest)
+//|
+//|      Encrypt the buffer from ``src`` into ``dest``.
+//|      For ECB mode, the buffers must be 16 bytes long.  For CBC mode, the
+//|      buffers must be a multiple of 16 bytes, and must be equal length.  For
+//|      CTX mode, there are no restrictions.
+//|
+STATIC mp_obj_t aesio_aes_encrypt_into(mp_obj_t aesio_obj, mp_obj_t src,
+                                       mp_obj_t dest) {
+  if (!MP_OBJ_IS_TYPE(aesio_obj, &aesio_aes_type)) {
+    mp_raise_TypeError_varg(translate("Expected a %q"), aesio_aes_type.name);
+  }
+  // Convert parameters into expected types.
+  aesio_aes_obj_t *aes = MP_OBJ_TO_PTR(aesio_obj);
+
+  mp_buffer_info_t srcbufinfo, destbufinfo;
+  mp_get_buffer_raise(src, &srcbufinfo, MP_BUFFER_READ);
+  mp_get_buffer_raise(dest, &destbufinfo, MP_BUFFER_READ);
+  validate_length(aes, srcbufinfo.len, destbufinfo.len);
+
+  memcpy(destbufinfo.buf, srcbufinfo.buf, srcbufinfo.len);
+
+  common_hal_aesio_aes_encrypt(aes, (uint8_t *)destbufinfo.buf,
+                               destbufinfo.len);
+  return mp_const_none;
+}
+
+STATIC MP_DEFINE_CONST_FUN_OBJ_3(aesio_aes_encrypt_into_obj,
+                                 aesio_aes_encrypt_into);
+
+//|   .. method:: decrypt_into(src, dest)
+//|
+//|      Decrypt the buffer from ``src`` into ``dest``.
+//|      For ECB mode, the buffers must be 16 bytes long.  For CBC mode, the
+//|      buffers must be a multiple of 16 bytes, and must be equal length.  For
+//|      CTX mode, there are no restrictions.
+//|
+STATIC mp_obj_t aesio_aes_decrypt_into(mp_obj_t aesio_obj, mp_obj_t src,
+                                       mp_obj_t dest) {
+  if (!MP_OBJ_IS_TYPE(aesio_obj, &aesio_aes_type)) {
+    mp_raise_TypeError_varg(translate("Expected a %q"), aesio_aes_type.name);
+  }
+  // Convert parameters into expected types.
+  aesio_aes_obj_t *aes = MP_OBJ_TO_PTR(aesio_obj);
+
+  mp_buffer_info_t srcbufinfo, destbufinfo;
+  mp_get_buffer_raise(src, &srcbufinfo, MP_BUFFER_READ);
+  mp_get_buffer_raise(dest, &destbufinfo, MP_BUFFER_READ);
+  validate_length(aes, srcbufinfo.len, destbufinfo.len);
+
+  memcpy(destbufinfo.buf, srcbufinfo.buf, srcbufinfo.len);
+
+  common_hal_aesio_aes_decrypt(aes, (uint8_t *)destbufinfo.buf,
+                               destbufinfo.len);
+  return mp_const_none;
+}
+
+STATIC MP_DEFINE_CONST_FUN_OBJ_3(aesio_aes_decrypt_into_obj,
+                                 aesio_aes_decrypt_into);
+
+STATIC mp_obj_t aesio_aes_get_mode(mp_obj_t aesio_obj) {
+  if (!MP_OBJ_IS_TYPE(aesio_obj, &aesio_aes_type)) {
+    mp_raise_TypeError_varg(translate("Expected a %q"), aesio_aes_type.name);
+  }
+  aesio_aes_obj_t *self = MP_OBJ_TO_PTR(aesio_obj);
+  return MP_OBJ_NEW_SMALL_INT(self->mode);
+}
+MP_DEFINE_CONST_FUN_OBJ_1(aesio_aes_get_mode_obj, aesio_aes_get_mode);
+
+STATIC mp_obj_t aesio_aes_set_mode(mp_obj_t aesio_obj, mp_obj_t mode_obj) {
+  if (!MP_OBJ_IS_TYPE(aesio_obj, &aesio_aes_type)) {
+    mp_raise_TypeError_varg(translate("Expected a %q"), aesio_aes_type.name);
+  }
+  aesio_aes_obj_t *self = MP_OBJ_TO_PTR(aesio_obj);
+
+  int mode = mp_obj_get_int(mode_obj);
+  switch (mode) {
+  case AES_MODE_CBC:
+  case AES_MODE_ECB:
+  case AES_MODE_CTR:
+    break;
+  default:
+    mp_raise_TypeError(translate("Requested AES mode is unsupported"));
+  }
+
+  common_hal_aesio_aes_set_mode(self, mode);
+  return mp_const_none;
+}
+MP_DEFINE_CONST_FUN_OBJ_2(aesio_aes_set_mode_obj, aesio_aes_set_mode);
+
+const mp_obj_property_t aesio_aes_mode_obj = {
+    .base.type = &mp_type_property,
+    .proxy = {
+      (mp_obj_t)&aesio_aes_get_mode_obj,
+      (mp_obj_t)&aesio_aes_set_mode_obj,
+      (mp_obj_t)&mp_const_none_obj
+    },
+};
+
+STATIC const mp_rom_map_elem_t aesio_locals_dict_table[] = {
+    // Methods
+    {MP_ROM_QSTR(MP_QSTR___name__), MP_OBJ_NEW_QSTR(MP_QSTR_AES)},
+    {MP_ROM_QSTR(MP_QSTR_encrypt_into), (mp_obj_t)&aesio_aes_encrypt_into_obj},
+    {MP_ROM_QSTR(MP_QSTR_decrypt_into), (mp_obj_t)&aesio_aes_decrypt_into_obj},
+    {MP_ROM_QSTR(MP_QSTR_rekey), (mp_obj_t)&aesio_aes_rekey_obj},
+    {MP_ROM_QSTR(MP_QSTR_mode), (mp_obj_t)&aesio_aes_mode_obj},
+};
+STATIC MP_DEFINE_CONST_DICT(aesio_locals_dict, aesio_locals_dict_table);
+
+const mp_obj_type_t aesio_aes_type = {
+    {&mp_type_type},
+    .name = MP_QSTR_AES,
+    .make_new = aesio_aes_make_new,
+    .locals_dict = (mp_obj_dict_t *)&aesio_locals_dict,
+};