1 files changed, 17 insertions, 24 deletions
diff --git a/python/lilliput.py b/python/lilliput.py
index f9c1b09..90a0ed1 100644
--- a/python/lilliput.py
+++ b/python/lilliput.py
@@ -1,40 +1,33 @@
+from enum import Enum
+
 import lilliput_ae_1
 import lilliput_ae_2
 from constants import NONCE_BYTES
 
 
-def _checkInputs(key, nonce):
+class LilliputAeMode(Enum):
+    lilliput_1 = lilliput_ae_1
+    lilliput_2 = lilliput_ae_2
+
+
+def _checkInputs(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 len(nonce) != NONCE_BYTES:
         raise ValueError('nonce must be {}-byte long'.format(NONCE_BYTES))
 
 
-def mainEnc(plaintext, adata, key, nonce, mode):
-    _checkInputs(key, nonce)
-
-    A = adata
-    M = plaintext
-    N = nonce
-
-    if(mode == 1) :
-        return lilliput_ae_1.OCB3Enc(A, M, N, key)
-    if(mode == 2) :
-        return lilliput_ae_2.SCT2Enc(A, M, N, key)
-
-
-def mainDec(ciphertext, tag, adata, key, nonce, mode):
-    _checkInputs(key, nonce)
+def encrypt(plaintext, adata, key, nonce, mode):
+    _checkInputs(key, mode, nonce)
+    return mode.value.encrypt(adata, plaintext, nonce, key)
 
-    A = adata
-    C = ciphertext
-    N = nonce
-    tag = list(tag)
 
-    if(mode == 1) :
-        return lilliput_ae_1.OCB3Dec(A, C, N, tag, key)
-    if(mode == 2) :
-        return lilliput_ae_2.SCT2Dec(A, C, N, tag, key)
+def decrypt(ciphertext, tag, adata, key, nonce, mode):
+    _checkInputs(key, mode, nonce)
+    return mode.value.decrypt(adata, ciphertext, nonce, tag, key)