Merge pull request #24 from ZhaoQi99/feat/add-license

support encrypt with license

Author: ZhaoQi99

README.md

@@ -1,5 +1,7 @@
 # pyencrypt-pye
+
 ---
+
 encrypt python source code and import module dynamically.
 
 ```
@@ -16,10 +18,12 @@ encrypt python source code and import module dynamically.
 ```
 
 ## Install
+
 ```bash
 pip install git+https://github.com/ZhaoQi99/pyencrypt-pye.git
 ✨🍰✨
 ```
+
 ## Usage
 
 ```shell
@@ -33,23 +37,45 @@ Commands:
   decrypt   Decrypt encrypted pye file
   encrypt   Encrypt your python code
   generate  Generate loader file using specified key
+  license   Generate license file using specified key
 ```
 
 ### Encrypt
+
 ```shell
 ~$ pyencrypt encrypt -h
 Usage: pyencrypt encrypt [OPTIONS] PATHNAME
 
   Encrypt your python code
 
 Options:
-  -i, --in-place  make changes to files in place
-  -k, --key TEXT  Your encryption key.If you don‘t specify key, pyencrypt will
-                  generate encryption key randomly.
-  --yes           Automatically answer yes for confirm questions.
-  -h, --help      Show this message and exit.
+  -i, --in-place                  make changes to files in place
+  -k, --key TEXT                  Your encryption key.If you don‘t specify
+                                  key, pyencrypt will generate encryption key
+                                  randomly.
+  --with-license                  Add license to encrypted file
+  -m, --bind-mac TEXT             Bind mac address to encrypted file
+  -4, --bind-ipv4 TEXT            Bind ipv4 address to encrypted file
+  -b, --before [%Y-%m-%dT%H:%M:%S %z|%Y-%m-%d %H:%M:%S|%Y-%m-%d]
+                                  License is invalid before this date.
+  -a, --after [%Y-%m-%dT%H:%M:%S %z|%Y-%m-%d %H:%M:%S|%Y-%m-%d]
+                                  License is invalid after this date.
+  -y, --yes                       Automatically answer yes for confirm
+                                  questions.
+  -h, --help                      Show this message and exit.
+```
+
+### Entry File
+
+In your entry file, you must import `loader` firstly, and then you can import encrypted modules as usual.
+
+```python
+import loader
+from test import *
 ```
+
 ### Decrypt
+
 ```shell
 ~$ pyencrypt decrypt -h
 Usage: pyencrypt decrypt [OPTIONS] PATHNAME
@@ -61,6 +87,7 @@ Options:
   -k, --key TEXT  Your encryption key.  [required]
   -h, --help      Show this message and exit.
 ```
+
 ### Generate
 
 ```shell
@@ -74,26 +101,63 @@ Options:
   -h, --help      Show this message and exit.
 ```
 
+### License
+
+pyencrypt's loader will search for the license file in the following manner:
+
+1. `~/.licenses/license.lic` file in your home directory.
+
+2. `licenses/license.lic` file in `loader` file's directory.
+
+3. `licenses/license.lic` file in the current working directory.
+
+```shell
+   ~$ pyencrypt license -h
+   Usage: pyencrypt license [OPTIONS]
+
+   Generate license file  using specified key
+
+Options:
+  -h, --help                      Show this message and exit.
+  -k, --key TEXT                  Your encryption key.  [required]
+  -m, --bind-mac TEXT             Your mac address.
+  -4, --bind-ipv4 TEXT            Your ipv4 address.
+  -b, --before [%Y-%m-%dT%H:%M:%S %z|%Y-%m-%d %H:%M:%S|%Y-%m-%d]
+                                  License is invalid before this date.
+  -a, --after [%Y-%m-%dT%H:%M:%S %z|%Y-%m-%d %H:%M:%S|%Y-%m-%d]
+                                  License is invalid after this date.
+```
+
 ## Example
+
 ### Encrypt
+
 ```shell
 ~$ pyencrypt encrypt --in-place  -y test.py
 ~$ pyencrypt encrypt test/
+~$ pyencrypt encrypt test.py -y --with-license\
+    --before="2000-01-01T00:00:00 +0800" --after="2030-01-01T00:00:00 +0800"\
+    --bind-mac="AC:DE:48:00:11:22" --bind-ipv4="192.168.0.1"
 ```
 
 ### Decrypt
+
 ```shell
 ~$ pyencrypt decrypt -k xxx test.pye
 ```
+
 ### Generate
+
 ```shell
 ~$ pyencrypt generate -k xxx
 ```
 
-### Entry File
-```python
-import loader
-from test import *
+### License
+
+```shell
+~$ pyencrypt license -k xxx\
+    --before="2000-01-01T00:00:00 +0800" --after="2030-01-01T00:00:00 +0800"\
+    --bind-mac="AC:DE:48:00:11:22" --bind-ipv4="192.168.0.1"
 ```
 
 ## Development
@@ -104,7 +168,11 @@ from test import *
 yapf --recursive -i pyencrypt 
 isort pyencrypt
 ```
+
 ## License
+
 [GNU General Public License v3.0](https://github.com/ZhaoQi99/pyencrypt-pye/blob/main/LICENSE)
+
 ## Author
+
 * Qi Zhao([[email protected]](mailto:[email protected]))

pyencrypt/cli.py

@@ -5,11 +5,12 @@
 from pathlib import Path
 
 import click
+
 from pyencrypt import __description__, __version__
 from pyencrypt.decrypt import decrypt_file
-from pyencrypt.encrypt import (can_encrypt, encrypt_file, encrypt_key,
-                               generate_so_file)
+from pyencrypt.encrypt import (can_encrypt, encrypt_file, encrypt_key, generate_so_file)
 from pyencrypt.generate import generate_aes_key
+from pyencrypt.license import generate_license_file, MIN_DATETIME, MAX_DATETIME
 
 VERSION = f"""\
                                                       _
@@ -33,16 +34,18 @@
 PYTHON_MAJOR, PYTHON_MINOR = sys.version_info[:2]
 LAODER_FILE_NAME = click.style(
     "encrypted/loader.cpython-{major}{minor}{abi}-{platform}.so".format(
-        major=PYTHON_MAJOR,
-        minor=PYTHON_MINOR,
-        abi=sys.abiflags,
-        platform=sys.platform),
+        major=PYTHON_MAJOR, minor=PYTHON_MINOR, abi=sys.abiflags, platform=sys.platform
+    ),
     blink=True,
-    fg='blue')
+    fg='blue'
+)
+LICENSE_FILE_NAME = click.style("license.lic", blink=True, fg='blue')
 
 SUCCESS_ANSI = click.style('successfully', fg='green')
 INVALID_KEY_MSG = click.style('Your encryption key is invalid.', fg='red')
 
+INVALID_EXPIRED_MSG = 'Expired before date must be less than expired after date.'
+
 FINISH_ENCRYPT_MSG = f"""
 Encryption completed {SUCCESS_ANSI}.
 Please copy {LAODER_FILE_NAME} into your encrypted directory.
@@ -54,10 +57,16 @@
 Decryption completed {SUCCESS_ANSI}. Your origin source code has be put: %s
 """
 
-FINISH_GENERATE_MSG = f"""
+FINISH_GENERATE_LOADER_MSG = f"""
 Generate loader file {SUCCESS_ANSI}. Your loader file is located in {LAODER_FILE_NAME}
 """
 
+FINISH_GENERATE_LICENSE_MSG = f"""
+Generate license file {SUCCESS_ANSI}. Your license file is located in {LICENSE_FILE_NAME}
+"""
+
+DATETIME_FORMATS = ['%Y-%m-%dT%H:%M:%S %z', '%Y-%m-%d %H:%M:%S', '%Y-%m-%d']
+
 
 def _check_key(key: str) -> bool:
     return not (len(key) % 4 or len(base64.b64decode(key)) % 16)
@@ -72,33 +81,26 @@ def cli():
 
 @cli.command(name='encrypt')
 @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',
-              default=None,
-              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.option('-i', '--in-place', 'replace', default=False, help='make changes to files in place', is_flag=True)
+@click.option('-k', '--key', default=None, help=KEY_OPTION_HELP, type=click.STRING)
+@click.option('--with-license', default=False, help='Add license to encrypted file', is_flag=True)
+@click.option('-m', '--bind-mac', 'mac', default=None, help='Bind mac address to encrypted file', type=click.STRING)
+@click.option('-4', '--bind-ipv4', 'ipv4', default=None, help='Bind ipv4 address to encrypted file', type=click.STRING)
+@click.option('-b', '--before', default=MAX_DATETIME, help='License is invalid before this date.', type=click.DateTime(formats=DATETIME_FORMATS))
+@click.option('-a', '--after', default=MIN_DATETIME, help='License is invalid after this date.', type=click.DateTime(formats=DATETIME_FORMATS))
+@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, replace, key):
+def encrypt_command(ctx, pathname, replace, key, with_license, mac, ipv4, before, after):
     """Encrypt your python code"""
     if key is not None and not _check_key(key):
         ctx.fail(INVALID_KEY_MSG)
     if key is None:
         key = generate_aes_key().decode()
-        click.echo(
-            f'Your randomly encryption 🔑 is {click.style(key,underline=True, fg="yellow")}'
-        )
+        click.echo(f'Your randomly encryption 🔑 is {click.style(key,underline=True, fg="yellow")}')
+
+    if before > after:
+        ctx.fail(INVALID_EXPIRED_MSG)
 
     path = Path(pathname)
 
@@ -126,22 +128,16 @@ def encrypt_command(ctx, pathname, replace, key):
 
     cipher_key, d, n = encrypt_key(key.encode())  # 需要放进导入器中
     generate_so_file(cipher_key, d, n)
+    if with_license is True:
+        generate_license_file(key, Path(os.getcwd()), after, before, mac, ipv4)
+        ctx.echo(FINISH_GENERATE_LICENSE_MSG)
     click.echo(FINISH_ENCRYPT_MSG)
 
 
 @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.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, replace, key):
@@ -177,20 +173,35 @@ def decrypt_command(ctx, pathname, replace, key):
 
 
 @cli.command(name='generate')
-@click.option('-k',
-              '--key',
-              required=True,
-              help='Your encryption key.',
-              type=click.STRING)
+@click.option('-k', '--key', required=True, help='Your encryption key.', type=click.STRING)
 @click.help_option('-h', '--help')
 @click.pass_context
 def generate_loader(ctx, key):
     """Generate loader file using specified key"""
     if not _check_key(key):
         ctx.fail(INVALID_KEY_MSG)
     cipher_key, d, n = encrypt_key(key.encode())
-    generate_so_file(cipher_key, d, n)
-    click.echo(FINISH_GENERATE_MSG)
+    generate_so_file(cipher_key, d, n, Path(os.getcwd()))
+    click.echo(FINISH_GENERATE_LOADER_MSG)
+
+
+@cli.command(name='license')
+@click.help_option('-h', '--help')
+@click.option('-k', '--key', required=True, help='Your encryption key.', type=click.STRING)
+@click.option('-m', '--bind-mac', help='Your mac address.', type=click.STRING)
+@click.option('-4', '--bind-ipv4', help='Your ipv4 address.', type=click.STRING)
+@click.option('-b', '--before', default=MAX_DATETIME, help='License is invalid before this date.', type=click.DateTime(formats=DATETIME_FORMATS))
+@click.option('-a', '--after', default=MIN_DATETIME, help='License is invalid after this date.', type=click.DateTime(formats=DATETIME_FORMATS))
+@click.pass_context
+def generate_license(ctx, key, mac, ipv4, before, after):
+    """Generate license file using specified key"""
+    if not _check_key(key):
+        ctx.fail(INVALID_KEY_MSG)
+    if before > after:
+        ctx.fail(INVALID_EXPIRED_MSG)
+
+    generate_license_file(key, Path(os.getcwd()), after, before, mac, ipv4)
+    click.echo(FINISH_GENERATE_LICENSE_MSG)
 
 
 if __name__ == '__main__':

pyencrypt/encrypt.py

@@ -1,4 +1,5 @@
 import os
+import re
 import subprocess
 from pathlib import Path
 
@@ -10,6 +11,8 @@
     '__init__.py',
 ]
 
+REMOVE_SELF_IMPORT = re.compile(r'^from pyencrypt\.[\s\S]*?$', re.MULTILINE)
+
 
 def _encrypt_file(
     data: bytes,
@@ -44,23 +47,22 @@ 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, base_dir: Path = None):
+def generate_so_file(cipher_key: str, d: int, n: int, base_dir: Path = None, license: bool = False) -> Path:
     private_key = f'{n}O{d}'
     path = Path(os.path.abspath(__file__)).parent
 
     decrypt_source_ls = list()
-    need_import_files = ['ntt.py', 'aes.py', 'decrypt.py']
+    need_import_files = ['ntt.py', 'aes.py', 'decrypt.py', 'license.py']
     for file in need_import_files:
         file_path = path / file
-        decrypt_source_ls.append(
-            file_path.read_text().replace('from pyencrypt.ntt import intt',
-                                          '').replace('from pyencrypt.aes import aes_decrypt', '')
-        )
+        decrypt_source_ls.append(REMOVE_SELF_IMPORT.sub('',file_path.read_text()))
 
     loader_source_path = path / 'loader.py'
-    loader_source = loader_source_path.read_text().replace(
+    loader_source = REMOVE_SELF_IMPORT.sub('', loader_source_path.read_text()).replace(
         "__private_key = ''", f"__private_key = '{private_key}'", 1
-    ).replace("__cipher_key = ''", f"__cipher_key = '{cipher_key}'", 1).replace("from pyencrypt.decrypt import *", '')
+    ).replace("__cipher_key = ''", f"__cipher_key = '{cipher_key}'", 1).replace(
+        'license = None', f'license = {license}', 1
+    )
 
     if base_dir is None:
         base_dir = Path(os.getcwd())
@@ -100,9 +102,7 @@ def generate_so_file(cipher_key: str, d: int, n: int, base_dir: Path = None):
         script_args=['build_ext', '--build-lib', temp_dir.as_posix()],
         cmdclass={'build_ext': build_ext},
     )
-
-    return True
-
+    return list(temp_dir.glob('loader.cpython-*-*.so'))[0].absolute()
 
 def encrypt_file(path: Path, key: str, delete_origin: bool = False, new_path: Path = None):
     if not can_encrypt(path):