1 files changed, 12 insertions, 14 deletions
diff --git a/python/lilliput/__init__.py b/python/lilliput/__init__.py
index 43179f8..5fbc0de 100644
--- a/python/lilliput/__init__.py
+++ b/python/lilliput/__init__.py
@@ -1,33 +1,31 @@
-from enum import Enum
-
 from . import lilliput_ae_1
 from . import lilliput_ae_2
 from .constants import NONCE_BYTES
 
 
-class LilliputAeMode(Enum):
-    lilliput_1 = lilliput_ae_1
-    lilliput_2 = lilliput_ae_2
+_AE_MODES = {
+    1: lilliput_ae_1,
+    2: lilliput_ae_2
+}
 
 
-def _checkInputs(key, mode, nonce):
+def _check_inputs(key, mode, nonce):
     valid_key_lengths = (128, 192, 256)
-
     if len(key)*8 not in valid_key_lengths:
         raise ValueError('invalid key size: {} not in {}'.format(len(key)*8, valid_key_lengths))
 
-    if mode.name not in LilliputAeMode.__members__:
-        raise ValueError('invalid mode: use a member of the LilliputAeMode enumeration')
+    if mode not in _AE_MODES:
+        raise ValueError('invalid mode: {} not in {}'.format(mode, tuple(_AE_MODES)))
 
     if len(nonce) != NONCE_BYTES:
-        raise ValueError('nonce must be {}-byte long'.format(NONCE_BYTES))
+        raise ValueError('invalid nonce size: expecting {}, have {}'.format(NONCE_BYTES, len(nonce)))
 
 
 def encrypt(plaintext, adata, key, nonce, mode):
-    _checkInputs(key, mode, nonce)
-    return mode.value.encrypt(adata, plaintext, nonce, key)
+    _check_inputs(key, mode, nonce)
+    return _AE_MODES[mode].encrypt(adata, plaintext, nonce, key)
 
 
 def decrypt(ciphertext, tag, adata, key, nonce, mode):
-    _checkInputs(key, mode, nonce)
-    return mode.value.decrypt(adata, ciphertext, nonce, tag, key)
+    _check_inputs(key, mode, nonce)
+    return _AE_MODES[mode].decrypt(adata, ciphertext, nonce, tag, key)