update AES module (#11)

* update AES module

* allow empty config files

Co-authored-by: QPod0 <[email protected]>

Author: haobibo

src/aloha/config/hocon.py

@@ -1,3 +1,4 @@
+from attrdict import AttrDict
 from pyhocon import ConfigFactory
 
 
@@ -14,4 +15,4 @@ def load_config_from_hocon_files(config_files: list, base_dir: str):
     f = '\n'.join(s)
 
     config = ConfigFactory.parse_string(content=f, basedir=base_dir).as_plain_ordered_dict()
-    return config
+    return AttrDict(config)

src/aloha/config/paths.py

@@ -1,6 +1,7 @@
 __all__ = ('get_resource_dir', 'get_config_dir', 'get_current_module_dir', 'get_project_base_dir', 'path_join')
 
 import os
+import warnings
 
 
 def path_join(*args) -> str:
@@ -50,10 +51,12 @@ def get_config_files() -> list:
     for f in files:
         file = get_config_dir(f)
         if not os.path.exists(file):
-            raise RuntimeError('Config file [%s] does not exists!' % file)
+            warnings.warn('Expecting config file [%s] but it does not exists!' % file)
         else:
             print('  ---> Loading config file [%s]' % file)
-        ret.append(os.path.expandvars(f))
+            ret.append(os.path.expandvars(f))
+    if len(ret) == 0:
+        warnings.warn('No config files set properly, EMPTY config will be used!')
     return ret
 
 

src/aloha/encrypt/aes.py

@@ -1,16 +1,16 @@
 import base64
 import binascii
-from typing import Union
+from typing import Union, Callable, Optional
 
 from Crypto.Cipher import AES
 from Crypto.Random import get_random_bytes
 from Crypto.Util.Padding import pad, unpad
 
-_AES_CIPHER_METHODS = {  # FULL_CIPHER_NAME: (dict_params,)
-    "AES/ECB/PKCS5Padding": {'mode': AES.MODE_ECB},
-    # "AES/ECB/NoPadding": {'mode': AES.MODE_ECB},
-    # "AES/CBC/PKCS5Padding": {'mode': AES.MODE_CBC, 'iv': b'0000000000000000'},
-    # "AES/CBC/NoPadding": {'mode': AES.MODE_CBC, 'iv': b'0000000000000000'},
+_AES_CIPHER_METHODS = {  # FULL_CIPHER_NAME: (dict_params, pad_style)
+    "AES/ECB/PKCS5Padding": ({'mode': AES.MODE_ECB}, 'pkcs7'),
+    "AES/ECB/NoPadding": ({'mode': AES.MODE_ECB}, 'pkcs7'),
+    "AES/CBC/PKCS7Padding": ({'mode': AES.MODE_CBC, 'iv': b'0000000000000000'}, 'pkcs7'),
+    "AES/CBC/NoPadding": ({'mode': AES.MODE_CBC, 'iv': b'0000000000000000'}, 'x923'),
 }
 
 
@@ -23,6 +23,8 @@ def _generate_key(key_size: int, method='const') -> bytes:
 
 
 class AesEncryptor:
+    supported_cipher_methods = _AES_CIPHER_METHODS
+
     def __init__(self, key: Union[str, bytes] = None, key_size: int = 16, cipher_name: str = 'AES/ECB/PKCS5Padding'):
         _key = key
         if key is None:
@@ -37,10 +39,13 @@ def __init__(self, key: Union[str, bytes] = None, key_size: int = 16, cipher_nam
         # https://pycryptodome.readthedocs.io/en/latest/src/util/util.html
         self.cipher_name = cipher_name
 
-    def encrypt(self, text: str, output_format='hex') -> Union[str, bytes]:
-        padded = pad(text.encode(), block_size=self.block_size)
+    def encrypt(self, text: str, output_format='hex', func_pad: Optional[Callable] = None) -> Union[str, bytes]:
+        dict_params, pad_style = _AES_CIPHER_METHODS.get(self.cipher_name)
+        if not callable(func_pad):
+            func_pad = lambda x: pad(data, block_size=self.block_size, style=pad_style)
 
-        dict_params = _AES_CIPHER_METHODS.get(self.cipher_name)
+        data = text.encode()
+        padded = func_pad(data)
         cipher = AES.new(key=self.key_aes, **dict_params)
         bytes_crypt = cipher.encrypt(padded)
 
@@ -54,7 +59,7 @@ def encrypt(self, text: str, output_format='hex') -> Union[str, bytes]:
             raise ValueError('Unknown output_type [%s]' % output_format)
         return crypt
 
-    def decrypt(self, text: Union[str, bytes], input_format: str = 'hex') -> Union[str, bytes]:
+    def decrypt(self, text: Union[str, bytes], input_format: str = 'hex', func_unpad: Optional[Callable] = None) -> Union[str, bytes]:
         text += (len(text) % 4) * '='
         if input_format == 'hex':
             crypt = binascii.a2b_hex(text)
@@ -64,10 +69,12 @@ def decrypt(self, text: Union[str, bytes], input_format: str = 'hex') -> Union[s
             crypt = text
         else:
             raise ValueError('Unknown output_type [%s]' % input_format)
-        dict_params = _AES_CIPHER_METHODS.get(self.cipher_name)
+        dict_params, pad_style = _AES_CIPHER_METHODS.get(self.cipher_name)
         cipher = AES.new(key=self.key_aes, **dict_params)
         data = cipher.decrypt(crypt)
-        data = unpad(data, block_size=self.block_size)
+        if not callable(func_unpad):
+            func_unpad = lambda x: unpad(x, block_size=self.block_size, style=pad_style)
+        data = func_unpad(data)
         return data.decode('UTF-8')
 
 

src/aloha/encrypt/rsa.py

@@ -21,6 +21,7 @@
 class RsaEncryptor:
     _dict_cache_cipher = {}
     _dict_cache_decipher = {}
+    supported_cipher_methods = _RSA_CIPHER_METHODS
 
     # ref: https://cryptobook.nakov.com/asymmetric-key-ciphers/rsa-encrypt-decrypt-examples
     def __init__(self, key_private: str = None, key_public: str = None, cipher_name: str = 'RSA/ECB/PKCS1Padding'):

src/setup.py

@@ -31,9 +31,9 @@
     zip_safe=False,
     install_requires=[],
     extras_require={
-        'base': ['pyhocon', 'pycryptodome'],
+        'base': ['attrdict3', 'pyhocon', 'pycryptodome'],
         'build': ['Cython'],
-        'service': ['requests', 'tornado', 'psutil'],
+        'service': ['requests', 'tornado', 'psutil', 'pyjwt'],
         'db': ['sqlalchemy', 'psycopg2-binary', 'pymysql', 'elasticsearch', 'pymongo', 'redis>4.2.0'],
         'stream': ['confluent_kafka'],
         'data': ['pandas'],