From dd78f83576adde648c7d15793bad696a039db621 Mon Sep 17 00:00:00 2001
From: Pavol Rusnak <stick@gk2.sk>
Date: Fri, 23 Sep 2016 18:33:10 +0200
Subject: [PATCH] update docu strings in aes

---
 extmod/modtrezorcrypto/modtrezorcrypto-aes.h |  7 ---
 src/trezor/crypto/aes.py                     | 50 ++++++++++++++++----
 2 files changed, 40 insertions(+), 17 deletions(-)

diff --git a/extmod/modtrezorcrypto/modtrezorcrypto-aes.h b/extmod/modtrezorcrypto/modtrezorcrypto-aes.h
index 4814972e..cb6cea33 100644
--- a/extmod/modtrezorcrypto/modtrezorcrypto-aes.h
+++ b/extmod/modtrezorcrypto/modtrezorcrypto-aes.h
@@ -9,9 +9,6 @@
 
 #include "trezor-crypto/aes.h"
 
-/*
-*/
-
 typedef struct _mp_obj_AES_t {
     mp_obj_base_t base;
     union {
@@ -23,10 +20,6 @@ typedef struct _mp_obj_AES_t {
     uint8_t ctr[AES_BLOCK_SIZE];
 } mp_obj_AES_t;
 
-/// def trezor.crypto.aes.AES(mode:int, key: bytes, iv: bytes=None) -> AES:
-///     '''
-///     Create AES context
-///     '''
 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);
diff --git a/src/trezor/crypto/aes.py b/src/trezor/crypto/aes.py
index 24c05f7d..2ae2fe98 100644
--- a/src/trezor/crypto/aes.py
+++ b/src/trezor/crypto/aes.py
@@ -1,31 +1,61 @@
 from TrezorCrypto import AES as _AES
 
-def AES_ECB_Encrypt(key):
+def AES_ECB_Encrypt(key: bytes):
+    '''
+    Create AES encryption context in ECB mode
+    '''
     return _AES(0x00, key)
 
-def AES_ECB_Decrypt(key):
+def AES_ECB_Decrypt(key: bytes):
+    '''
+    Create AES decryption context in ECB mode
+    '''
     return _AES(0x80, key)
 
-def AES_CBC_Encrypt(key, iv):
+def AES_CBC_Encrypt(key: bytes, iv: bytes):
+    '''
+    Create AES encryption context in CBC mode
+    '''
     return _AES(0x01, key, iv)
 
-def AES_CBC_Decrypt(key, iv):
+def AES_CBC_Decrypt(key: bytes, iv: bytes):
+    '''
+    Create AES decryption context in CBC mode
+    '''
     return _AES(0x81, key, iv)
 
-def AES_CFB_Encrypt(key, iv):
+def AES_CFB_Encrypt(key: bytes, iv: bytes):
+    '''
+    Create AES encryption context in CFB mode
+    '''
     return _AES(0x02, key, iv)
 
-def AES_CFB_Decrypt(key, iv):
+def AES_CFB_Decrypt(key: bytes, iv: bytes):
+    '''
+    Create AES decryption context in CFB mode
+    '''
     return _AES(0x82, key, iv)
 
-def AES_OFB_Encrypt(key, iv):
+def AES_OFB_Encrypt(key: bytes, iv: bytes):
+    '''
+    Create AES encryption context in OFB mode
+    '''
     return _AES(0x03, key, iv)
 
-def AES_OFB_Decrypt(key, iv):
+def AES_OFB_Decrypt(key: bytes, iv: bytes):
+    '''
+    Create AES decryption context in OFB mode
+    '''
     return _AES(0x83, key, iv)
 
-def AES_CTR_Encrypt(key):
+def AES_CTR_Encrypt(key: bytes):
+    '''
+    Create AES encryption context in CTR mode
+    '''
     return _AES(0x04, key)
 
-def AES_CTR_Decrypt(key):
+def AES_CTR_Decrypt(key: bytes):
+    '''
+    Create AES decryption context in CTR mode
+    '''
     return _AES(0x84, key)