From 473fe45da23f371d7beb6117c31922f9c9ee5942 Mon Sep 17 00:00:00 2001
From: Yonatan Goldschmidt <yon.goldschmidt@gmail.com>
Date: Fri, 15 Jun 2018 17:07:47 +0300
Subject: [PATCH] extmod/moducryptolib: Optionally export MODE_* constants to
 Python.

Allow including crypto consts based on compilation settings.  Disabled by
default to reduce code size; if one wants extra code readability, can
enable them.
---
 docs/library/ucryptolib.rst |  4 ++--
 extmod/moducryptolib.c      | 24 ++++++++++++++++++------
 py/mpconfig.h               |  4 ++++
 3 files changed, 24 insertions(+), 8 deletions(-)

diff --git a/docs/library/ucryptolib.rst b/docs/library/ucryptolib.rst
index 4b2c45f073..c9e0bb71f7 100644
--- a/docs/library/ucryptolib.rst
+++ b/docs/library/ucryptolib.rst
@@ -21,8 +21,8 @@ Classes
             * *key* is an encryption/decryption key (bytes-like).
             * *mode* is:
 
-                * 1 for Electronic Code Book (ECB).
-                * 2 for Cipher Block Chaining (CBC)
+                * ``1`` (or ``ucryptolib.MODE_ECB`` if it exists) for Electronic Code Book (ECB).
+                * ``2`` (or ``ucryptolib.MODE_CBC`` if it exists) for Cipher Block Chaining (CBC)
 
             * *IV* is an initialization vector for CBC mode.
 
diff --git a/extmod/moducryptolib.c b/extmod/moducryptolib.c
index 88b3447bbe..4f3a6b8e84 100644
--- a/extmod/moducryptolib.c
+++ b/extmod/moducryptolib.c
@@ -4,6 +4,7 @@
  * The MIT License (MIT)
  *
  * Copyright (c) 2017-2018 Paul Sokolovsky
+ * Copyright (c) 2018 Yonatan Goldschmidt
  *
  * Permission is hereby granted, free of charge, to any person obtaining a copy
  * of this software and associated documentation files (the "Software"), to deal
@@ -25,6 +26,7 @@
  */
 
 #include "py/mpconfig.h"
+
 #if MICROPY_PY_UCRYPTOLIB
 
 #include <assert.h>
@@ -38,17 +40,23 @@
 // of PEP 272 can be made with a simple wrapper which adds all the
 // needed boilerplate.
 
+// values follow PEP 272
+enum {
+    UCRYPTOLIB_MODE_MIN = 0,
+    UCRYPTOLIB_MODE_ECB,
+    UCRYPTOLIB_MODE_CBC,
+    UCRYPTOLIB_MODE_MAX,
+};
+
 #if MICROPY_SSL_AXTLS
 #include "lib/axtls/crypto/crypto.h"
 #endif
 
-#define MODE_ECB 1
-#define MODE_CBC 2
 
 typedef struct _mp_obj_aes_t {
     mp_obj_base_t base;
     AES_CTX ctx;
-    uint8_t mode: 7;
+    uint8_t block_mode: 7;
     uint8_t is_decrypt_key: 1;
 } mp_obj_aes_t;
 
@@ -57,10 +65,10 @@ STATIC mp_obj_t aes_make_new(const mp_obj_type_t *type, size_t n_args, size_t n_
     mp_obj_aes_t *o = m_new_obj(mp_obj_aes_t);
     o->base.type = type;
 
-    o->mode = mp_obj_get_int(args[1]);
+    o->block_mode = mp_obj_get_int(args[1]);
     o->is_decrypt_key = 0;
 
-    if (o->mode < MODE_ECB || o->mode > MODE_CBC) {
+    if (o->block_mode <= UCRYPTOLIB_MODE_MIN || o->block_mode >= UCRYPTOLIB_MODE_MAX) {
         mp_raise_ValueError("mode");
     }
 
@@ -123,7 +131,7 @@ STATIC mp_obj_t aes_process(size_t n_args, const mp_obj_t *args, bool encrypt) {
         self->is_decrypt_key = 1;
     }
 
-    if (self->mode == MODE_ECB) {
+    if (self->block_mode == UCRYPTOLIB_MODE_ECB) {
             uint8_t *in = in_bufinfo.buf, *out = out_buf_ptr;
             uint8_t *top = in + in_bufinfo.len;
             for (; in < top; in += 16, out += 16) {
@@ -183,6 +191,10 @@ STATIC const mp_obj_type_t aes_type = {
 STATIC const mp_rom_map_elem_t mp_module_ucryptolib_globals_table[] = {
     { MP_ROM_QSTR(MP_QSTR___name__), MP_ROM_QSTR(MP_QSTR_ucryptolib) },
     { MP_ROM_QSTR(MP_QSTR_aes), MP_ROM_PTR(&aes_type) },
+#if MICROPY_PY_UCRYPTOLIB_CONSTS
+    { MP_ROM_QSTR(MP_QSTR_MODE_ECB), MP_ROM_INT(UCRYPTOLIB_MODE_ECB) },
+    { MP_ROM_QSTR(MP_QSTR_MODE_CBC), MP_ROM_INT(UCRYPTOLIB_MODE_CBC) },
+#endif
 };
 
 STATIC MP_DEFINE_CONST_DICT(mp_module_ucryptolib_globals, mp_module_ucryptolib_globals_table);
diff --git a/py/mpconfig.h b/py/mpconfig.h
index 2fe3bdb5a1..8f202380da 100644
--- a/py/mpconfig.h
+++ b/py/mpconfig.h
@@ -1167,6 +1167,10 @@ typedef double mp_float_t;
 #define MICROPY_PY_UCRYPTOLIB (0)
 #endif
 
+#ifndef MICROPY_PY_UCRYPTOLIB_CONSTS
+#define MICROPY_PY_UCRYPTOLIB_CONSTS (0)
+#endif
+
 #ifndef MICROPY_PY_UBINASCII
 #define MICROPY_PY_UBINASCII (0)
 #endif