fix: 🐛 support append sys.path

Author: ZhaoQi99

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)
 
 
@@ -44,7 +44,7 @@ def encrypt_key(key: bytes) -> str:
     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
 
@@ -62,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)
@@ -97,6 +100,7 @@ def generate_so_file(cipher_key: str, d: int, n: int):
                          encoding='utf-8')
     if ret.returncode == 0:
         pass
+    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