Merge pull request #19 from ZhaoQi99/feat/add-unit-test

Feat/add unit test

Author: ZhaoQi99

dev-requirement.txt

@@ -1,3 +1,5 @@
 -r requirement.txt
 isort==5.10.1
-yapf==0.31.0
+yapf==0.31.0
+pytest==6.2.5
+ipython==7.16.3

pyencrypt/aes.py

@@ -848,21 +848,3 @@ def aes_decrypt(data: bytes, key: str) -> bytes:
         plain.append(AESModeOfOperationECB(key).decrypt(x))
     return strip_padding(b''.join(plain))
 
-
-if __name__ == '__main__':
-    plain = '你好!世界!'
-    key = 'tB9qW0YGlYIyBfTmuyQbm6AjPQa9gTKwO8j5H4IBf1A='
-    plain = add_padding(plain.encode())
-    print('plain', plain)
-    cipher = AESModeOfOperationECB(key).encrypt(plain)
-    print('cipher', cipher)
-    plain = strip_padding(AESModeOfOperationECB(key).decrypt(cipher))
-    print(plain)
-    print(plain.decode())
-
-    with open('generate.py', 'rb') as f:
-        data = f.read()
-    cipher = aes_encrypt(data, key)
-    print(cipher)
-    plain = aes_decrypt(cipher, key)
-    print(plain.decode())

pyencrypt/decrypt.py

@@ -4,7 +4,7 @@
 from pyencrypt.ntt import intt
 
 
-def decrypt_key(cipher_key: str, n: int, d: int) -> str:
+def decrypt_key(cipher_key: str, d: int, n: int) -> str:
     plain_ls = list()
     for num in map(int, cipher_key.split('O')):
         plain_ls.append(pow(num, d, n))

pyencrypt/encrypt.py

@@ -14,7 +14,7 @@
 def _encrypt_file(
     data: bytes,
     key: bytes,
-) -> None:
+) -> bytes:
     return aes_encrypt(data, key)
 
 
@@ -32,14 +32,19 @@ def encrypt_key(key: bytes) -> str:
     ascii_ls = [ord(x) for x in key.decode()]
     numbers = generate_rsa_number(2048)
     e, n = numbers['e'], numbers['n']
+    # fill length to be a power of 2
+    length = len(ascii_ls)
+    if length & (length - 1) != 0:
+        length = 1 << length.bit_length()
+        ascii_ls = ascii_ls + [0] * (length - len(ascii_ls))
     cipher_ls = list()
     # ntt后再用RSA加密
     for num in ntt(ascii_ls):
         cipher_ls.append(pow(num, e, n))
     return 'O'.join(map(str, cipher_ls)), numbers['d'], numbers['n']
 
 
-def generate_so_file(cipher_key: str, d: int, n: int):
+def generate_so_file(cipher_key: str, d: int, n: int, base_dir: Path = None):
     private_key = f'{n}O{d}'
     path = Path(os.path.abspath(__file__)).parent
 
@@ -57,7 +62,10 @@ def generate_so_file(cipher_key: str, d: int, n: int):
         1).replace("__cipher_key = ''", f"__cipher_key = '{cipher_key}'",
                    1).replace("from pyencrypt.decrypt import *", '')
 
-    temp_dir = Path(os.getcwd()) / 'encrypted'
+    if base_dir is None:
+        base_dir = Path(os.getcwd())
+
+    temp_dir = base_dir / 'encrypted'
     temp_dir.mkdir(exist_ok=True)
     loader_file_path = temp_dir / 'loader.py'
     loader_file_path.touch(exist_ok=True)
@@ -70,7 +78,6 @@ def generate_so_file(cipher_key: str, d: int, n: int):
     loader_origin_file_path.touch(exist_ok=True)
     loader_origin_file_path.write_text(f"{decrypt_source}\n{loader_source}")
 
-    setup_file_path = Path(os.path.abspath(__file__)).parent / 'setup.py'
     args = [
         'pyminifier', '--obfuscate-classes', '--obfuscate-import-methods',
         '--replacement-length', '20', '-o',
@@ -81,17 +88,16 @@ def generate_so_file(cipher_key: str, d: int, n: int):
     if ret.returncode == 0:
         pass
 
-    args = [
-        'python',
-        setup_file_path.as_posix(), 'build_ext', '--build-lib',
-        temp_dir.as_posix()
-    ]
-    ret = subprocess.run(args,
-                         shell=False,
-                         stderr=subprocess.PIPE,
-                         encoding='utf-8')
-    if ret.returncode == 0:
-        pass
+    from setuptools import setup
+    from Cython.Build import cythonize
+    from Cython.Distutils import build_ext
+    setup(
+        ext_modules=cythonize(loader_file_path.as_posix(), language_level="3"),
+        script_args=['build_ext', '--build-lib', temp_dir.as_posix()],
+        cmdclass={'build_ext': build_ext},
+    )
+
+    return True
 
 
 def encrypt_file(path: Path,

pyencrypt/generate.py

@@ -4,15 +4,11 @@
 from Crypto.PublicKey import RSA
 
 
-def generate_aes_key(size: int = 32) -> str:
+def generate_aes_key(size: int = 32) -> bytes:
     return base64.b64encode(os.urandom(size))
 
 
-def generate_rsa_number(bits: int):
+def generate_rsa_number(bits: int)-> dict:
     r = RSA.generate(bits)
     return {'p': r.p, 'q': r.q, 'n': r.n, 'e': r.e, 'd': r.d}
 
-
-if __name__ == '__main__':
-    print(generate_aes_key(32))
-    print(generate_rsa_number(2048))

pyencrypt/loader.py

@@ -33,7 +33,7 @@ def get_data(self, path: _Path) -> bytes:
         try:
             __n, __d = self.__private_key.split('O', 1)
             return decrypt_file(
-                Path(path), decrypt_key(self.__cipher_key, int(__n), int(__d)))
+                Path(path), decrypt_key(self.__cipher_key, int(__d), int(__n)))
         except Exception:
             traceback.print_exc()
             return b''
@@ -51,7 +51,10 @@ def find_spec(self,
             else:
                 file_path = Path(path[0]) / f'{fullname.rsplit(".",1)[-1]}.pye'
         else:
-            file_path = Path(f'{fullname}.pye')
+            for p in sys.path:
+                file_path = Path(p) / f'{fullname}.pye'
+                if file_path.exists():
+                    break
         file_path = file_path.absolute().as_posix()
         if not os.path.exists(file_path):
             return None

pyencrypt/ntt.py

@@ -16,10 +16,10 @@ def bitreverse(x: list, length: int):
 
 
 def _ntt(arr: list, inverse=False):
-    length = 1
-    while length < len(arr):
-        length *= 2
-    x = arr + [0] * (length - len(arr))
+    length = len(arr)
+    if length & (length - 1) != 0:
+        raise ValueError("The length of input must be a power of 2.")
+    x = arr.copy()
     g = pow(G, (M - 1) // length, M)
     if inverse:
         g = pow(g, M - 2, M)

pyencrypt/setup.py

@@ -20,7 +20,7 @@ classifiers =
     Programming Language :: Python :: 3.9
     Programming Language :: Python :: Implementation :: CPython
     Environment :: Console
-license_file = LICENSE
+license_files = LICENSE
 description = A simple image converter for Python.
 long_description = file: README.md
 long_description_content_type = text/markdown
@@ -45,4 +45,11 @@ exclude =
 
 [options.entry_points]
 console_scripts =
-    pyencrypt = pyencrypt.cli:cli
+    pyencrypt = pyencrypt.cli:cli
+
+[yapf]
+based_on_style = pep8
+COLUMN_LIMIT = 119
+indent_width =  4
+coalesce_brackets = True
+dedent_closing_brackets = True

setup.cfg

@@ -0,0 +1 @@
+AES_KEY = b'tiovaZzCF/Hlx/8uw0PXCkiqxCXGosP/AKo6Kn1gECw='