perf: ⚡️ encrypt and decrypt command

Author: ZhaoQi99

README.md

@@ -57,6 +57,7 @@ Usage: pyencrypt decrypt [OPTIONS] PATHNAME
   Decrypt encrypted pye file
 
 Options:
+  -i, --in-place  make changes to files in place
   -k, --key TEXT  Your encryption key.  [required]
   -h, --help      Show this message and exit.
 ```

pyencrypt/cli.py

@@ -74,7 +74,7 @@ def cli():
 @click.argument('pathname', type=click.Path(exists=True, resolve_path=True))
 @click.option('-i',
               '--in-place',
-              'delete',
+              'replace',
               default=False,
               help='make changes to files in place',
               is_flag=True)
@@ -84,11 +84,13 @@ def cli():
               help=KEY_OPTION_HELP,
               type=click.STRING)
 @click.confirmation_option(
+    '-y',
+    '--yes',
     prompt='Are you sure you want to encrypt your python file?',
     help='Automatically answer yes for confirm questions.')
 @click.help_option('-h', '--help')
 @click.pass_context
-def encrypt_command(ctx, pathname, delete, key):
+def encrypt_command(ctx, pathname, replace, key):
     """Encrypt your python code"""
     if key is not None and not _check_key(key):
         ctx.fail(INVALID_KEY_MSG)
@@ -99,23 +101,26 @@ def encrypt_command(ctx, pathname, delete, key):
         )
 
     path = Path(pathname)
-    work_dir = Path(os.getcwd()) / 'encrypted' / 'src'
 
     if path.is_file():
-        new_path = Path(os.getcwd()) / path.with_suffix('.pye').name
-        encrypt_file(path, key, delete, new_path)
+        if replace:
+            new_path = path.with_suffix('.pye')
+        else:
+            new_path = Path(os.getcwd()) / path.with_suffix('.pye').name
+        encrypt_file(path, key, replace, new_path)
     elif path.is_dir():
-        work_dir.exists() and shutil.rmtree(work_dir)
-        shutil.copytree(path, work_dir)
-        files = set(path.glob('**/*.py')) - set(
-            path.glob(f'encrypted/**/*.py'))
+        if replace:
+            work_dir = path
+        else:
+            work_dir = Path(os.getcwd()) / 'encrypted' / path.name
+            work_dir.exists() and shutil.rmtree(work_dir)
+            shutil.copytree(path, work_dir)
+        files = set(work_dir.glob('**/*.py'))
         with click.progressbar(files, label='🔐 Encrypting') as bar:
             for file in bar:
+                new_path = file.with_suffix('.pye')
                 if can_encrypt(file):
-                    new_path = work_dir / file.relative_to(path)
-                    new_path.unlink()
-                    encrypt_file(file, key, delete,
-                                 new_path.with_suffix('.pye'))
+                    encrypt_file(file, key, True, new_path)
     else:
         raise Exception(f'{path} is not a valid path.')
 
@@ -126,33 +131,45 @@ def encrypt_command(ctx, pathname, delete, key):
 
 @cli.command(name='decrypt')
 @click.argument('pathname', type=click.Path(exists=True, resolve_path=True))
+@click.option('-i',
+              '--in-place',
+              'replace',
+              default=False,
+              help='make changes to files in place',
+              is_flag=True)
 @click.option('-k',
               '--key',
               required=True,
               help='Your encryption key.',
               type=click.STRING)
 @click.help_option('-h', '--help')
 @click.pass_context
-def decrypt_command(ctx, pathname, key):
+def decrypt_command(ctx, pathname, replace, key):
     """Decrypt encrypted pye file"""
     path = Path(pathname)
     if not _check_key(key):
         ctx.fail(INVALID_KEY_MSG)
 
     if path.is_file():
-        work_dir = Path(os.getcwd())
-        new_path = work_dir / path.with_suffix('.py').name
-        origin_data = decrypt_file(path, key, new_path)
+        if replace:
+            new_path = path.with_suffix('.py')
+        else:
+            new_path = Path(os.getcwd()) / path.with_suffix('.py').name
+        work_dir = new_path.parent
+        origin_data = decrypt_file(path, key, replace, new_path)
         print(origin_data.decode())
     elif path.is_dir():
-        work_dir = Path(os.getcwd()) / 'decrypted' / 'src'
-        work_dir.exists() and shutil.rmtree(work_dir)
-        shutil.copytree(path, work_dir)
-        files = list(path.glob('**/*.pye'))
+        if replace:
+            work_dir = path
+        else:
+            work_dir = Path(os.getcwd()) / 'decrypted' / path.name
+            work_dir.exists() and shutil.rmtree(work_dir)
+            shutil.copytree(path, work_dir)
+        files = list(work_dir.glob('**/*.pye'))
         with click.progressbar(files, label='🔓 Decrypting') as bar:
-            for file in files:
-                new_path = work_dir / file.relative_to(path)
-                decrypt_file(file, key, new_path)
+            for file in bar:
+                new_path = file.with_suffix('.py')
+                decrypt_file(file, key, True, new_path)
     else:
         raise Exception(f'{path} is not a valid path.')
 

pyencrypt/decrypt.py

@@ -1,5 +1,5 @@
 from pathlib import Path
-
+import os
 from pyencrypt.aes import aes_decrypt
 from pyencrypt.ntt import intt
 
@@ -16,7 +16,10 @@ def _decrypt_file(data: bytes, key: str) -> bytes:
     return aes_decrypt(data, key)
 
 
-def decrypt_file(path: Path, key: str, new_path: Path = None) -> bytes:
+def decrypt_file(path: Path,
+                 key: str,
+                 delete_origin: bool = False,
+                 new_path: Path = None) -> bytes:
     if path.suffix != '.pye':
         raise Exception(f"{path.name} can't be decrypted.")
     data = _decrypt_file(path.read_bytes(), key)
@@ -25,4 +28,6 @@ def decrypt_file(path: Path, key: str, new_path: Path = None) -> bytes:
             raise Exception("Origin file path must be py suffix.")
         new_path.touch(exist_ok=True)
         new_path.write_bytes(data)
+    if delete_origin:
+        os.remove(path)
     return data