Skip to content

Commit

Permalink
style: 💄 format code using black
Browse files Browse the repository at this point in the history
  • Loading branch information
@ZhaoQi99
ZhaoQi99 committed Mar 26, 2024
1 parent 2f825ad commit 4a0bc4b
Show file tree
Hide file tree
Showing 16 changed files with 659 additions and 406 deletions.
145 changes: 74 additions & 71 deletions pyencrypt/aes.py
View file

Large diffs are not rendered by default.

201 changes: 139 additions & 62 deletions pyencrypt/cli.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,11 +10,11 @@

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 MAX_DATETIME, MIN_DATETIME, generate_license_file

VERSION = fr"""
VERSION = rf"""
_
_ __ _ _ ___ _ __ ___ _ __ _ _ _ __ | |_
| '_ \| | | |/ _ \ '_ \ / __| '__| | | | '_ \| __|
Expand All @@ -33,18 +33,20 @@
"""

PYTHON_MAJOR, PYTHON_MINOR = sys.version_info[:2]
LOADER_FILE_NAME = click.style("encrypted/{}", blink=True, fg='blue')
LICENSE_FILE_NAME = click.style("license.lic", blink=True, fg='blue')
LOADER_FILE_NAME = click.style("encrypted/{}", blink=True, fg="blue")
LICENSE_FILE_NAME = click.style("license.lic", blink=True, fg="blue")

SUCCESS_ANSI = click.style('successfully', fg='green')
SUCCESS_ANSI = click.style("successfully", fg="green")

INVALID_KEY_MSG = click.style('Your encryption 🔑 is invalid.', fg='red')
INVALID_KEY_MSG = click.style("Your encryption 🔑 is invalid.", fg="red")

INVALID_MAC_MSG = click.style('{} is not a valid mac address.', fg='red')
INVALID_MAC_MSG = click.style("{} is not a valid mac address.", fg="red")

INVALID_IPV4_MSG = click.style('{} is not a valid ipv4 address.', fg='red')
INVALID_IPV4_MSG = click.style("{} is not a valid ipv4 address.", fg="red")

INVALID_DATETIME_MSG = click.style('Before date must be less than after date.', fg='red')
INVALID_DATETIME_MSG = click.style(
"Before date must be less than after date.", fg="red"
)

FINISH_ENCRYPT_MSG = f"""
Encryption completed {SUCCESS_ANSI}.
Expand All @@ -65,11 +67,11 @@
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']
DATETIME_FORMATS = ["%Y-%m-%dT%H:%M:%S %z", "%Y-%m-%d %H:%M:%S", "%Y-%m-%d"]


class KeyParamType(click.ParamType):
name = 'key'
name = "key"

def _check_key(self, key: str) -> bool:
return not (len(key) % 4 or len(base64.b64decode(key)) % 16)
Expand All @@ -81,15 +83,15 @@ def convert(self, value, param, ctx) -> str:
return value

def get_metavar(self, param):
return '🔑'
return "🔑"

def __repr__(self) -> str:
return "KEY"


class MacAddressParamType(click.ParamType):
name = 'mac_address'
pattern = re.compile(r'^([0-9a-fA-F]{2}[:-]){5}([0-9a-fA-F]{2})$')
name = "mac_address"
pattern = re.compile(r"^([0-9a-fA-F]{2}[:-]){5}([0-9a-fA-F]{2})$")

def convert(self, value, param, ctx) -> str:
value = click.STRING.convert(value, param, ctx)
Expand All @@ -98,14 +100,14 @@ def convert(self, value, param, ctx) -> str:
return value

def get_metavar(self, param):
return '01:23:45:67:89:AB'
return "01:23:45:67:89:AB"

def __repr__(self) -> str:
return "MacAddress"


class IPv4AddressParamType(click.ParamType):
name = 'ipv4_address'
name = "ipv4_address"

def convert(self, value, param, ctx) -> str:
value = click.STRING.convert(value, param, ctx)
Expand All @@ -115,7 +117,7 @@ def convert(self, value, param, ctx) -> str:
self.fail(INVALID_IPV4_MSG.format(value), param, ctx)

def get_metavar(self, param):
return '192.168.0.1'
return "192.168.0.1"

def __repr__(self) -> str:
return "Ipv4Address"
Expand All @@ -128,29 +130,75 @@ class CustomParamType:


@click.group()
@click.version_option(__version__, '--version', message=VERSION)
@click.help_option('-h', '--help')
@click.version_option(__version__, "--version", message=VERSION)
@click.help_option("-h", "--help")
def cli():
pass


@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=CustomParamType.KEY)
@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=CustomParamType.MAC_ADDR)
@click.option('-4', '--bind-ipv4', 'ipv4', default=None, help='Bind ipv4 address to encrypted file', type=CustomParamType.IPV4_ADDR)
@click.option('-b', '--before', default=MIN_DATETIME, help='License is invalid before this date.', type=click.DateTime(formats=DATETIME_FORMATS))
@click.option('-a', '--after', default=MAX_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')
@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=CustomParamType.KEY
)
@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=CustomParamType.MAC_ADDR,
)
@click.option(
"-4",
"--bind-ipv4",
"ipv4",
default=None,
help="Bind ipv4 address to encrypted file",
type=CustomParamType.IPV4_ADDR,
)
@click.option(
"-b",
"--before",
default=MIN_DATETIME,
help="License is invalid before this date.",
type=click.DateTime(formats=DATETIME_FORMATS),
)
@click.option(
"-a",
"--after",
default=MAX_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, with_license, mac, ipv4, before, after):
def encrypt_command(
ctx, pathname, replace, key, with_license, mac, ipv4, before, after
):
"""Encrypt your python code"""
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_DATETIME_MSG)
Expand All @@ -159,25 +207,25 @@ def encrypt_command(ctx, pathname, replace, key, with_license, mac, ipv4, before

if path.is_file():
if replace:
new_path = path.with_suffix('.pye')
new_path = path.with_suffix(".pye")
else:
new_path = Path(os.getcwd()) / path.with_suffix('.pye').name
new_path = Path(os.getcwd()) / path.with_suffix(".pye").name
encrypt_file(path, key, replace, new_path)
elif path.is_dir():
if replace:
work_dir = path
else:
work_dir = Path(os.getcwd()) / 'encrypted' / path.name
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:
files = set(work_dir.glob("**/*.py"))
with click.progressbar(files, label="🔐 Encrypting") as bar:
for file in bar:
new_path = file.with_suffix('.pye')
new_path = file.with_suffix(".pye")
if can_encrypt(file):
encrypt_file(file, key, True, new_path)
else:
raise Exception(f'{path} is not a valid path.')
raise Exception(f"{path} is not a valid path.")

cipher_key, d, n = encrypt_key(key.encode()) # 需要放进导入器中
loader_extension = generate_so_file(cipher_key, d, n, license=with_license)
Expand All @@ -187,45 +235,56 @@ def encrypt_command(ctx, pathname, replace, key, with_license, mac, ipv4, before
click.echo(FINISH_ENCRYPT_MSG.format(loader_extension.name))


@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=CustomParamType.KEY)
@click.help_option('-h', '--help')
@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=CustomParamType.KEY
)
@click.help_option("-h", "--help")
@click.pass_context
def decrypt_command(ctx, pathname, replace, key):
"""Decrypt encrypted pye file"""
path = Path(pathname)

if path.is_file():
if replace:
new_path = path.with_suffix('.py')
new_path = path.with_suffix(".py")
else:
new_path = Path(os.getcwd()) / path.with_suffix('.py').name
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():
if replace:
work_dir = path
else:
work_dir = Path(os.getcwd()) / 'decrypted' / path.name
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:
files = list(work_dir.glob("**/*.pye"))
with click.progressbar(files, label="🔓 Decrypting") as bar:
for file in bar:
new_path = file.with_suffix('.py')
new_path = file.with_suffix(".py")
decrypt_file(file, key, True, new_path)
else:
raise Exception(f'{path} is not a valid path.')
raise Exception(f"{path} is not a valid path.")

click.echo(FINISH_DECRYPT_MSG.format(work_dir=work_dir))


@cli.command(name='generate')
@click.option('-k', '--key', required=True, help='Your encryption key.', type=CustomParamType.KEY)
@click.help_option('-h', '--help')
@cli.command(name="generate")
@click.option(
"-k", "--key", required=True, help="Your encryption key.", type=CustomParamType.KEY
)
@click.help_option("-h", "--help")
@click.pass_context
def generate_loader(ctx, key):
"""Generate loader file using specified key"""
Expand All @@ -234,13 +293,31 @@ def generate_loader(ctx, key):
click.echo(FINISH_GENERATE_LOADER_MSG.format(loader_extension.name))


@cli.command(name='license')
@click.help_option('-h', '--help')
@click.option('-k', '--key', required=True, help='Your encryption key.', type=CustomParamType.KEY)
@click.option('-m', '--bind-mac', help='Your mac address.', type=CustomParamType.MAC_ADDR)
@click.option('-4', '--bind-ipv4', help='Your ipv4 address.', type=CustomParamType.IPV4_ADDR)
@click.option('-b', '--before', default=MIN_DATETIME, help='License is invalid before this date.', type=click.DateTime(formats=DATETIME_FORMATS))
@click.option('-a', '--after', default=MAX_DATETIME, help='License is invalid after this date.', type=click.DateTime(formats=DATETIME_FORMATS))
@cli.command(name="license")
@click.help_option("-h", "--help")
@click.option(
"-k", "--key", required=True, help="Your encryption key.", type=CustomParamType.KEY
)
@click.option(
"-m", "--bind-mac", help="Your mac address.", type=CustomParamType.MAC_ADDR
)
@click.option(
"-4", "--bind-ipv4", help="Your ipv4 address.", type=CustomParamType.IPV4_ADDR
)
@click.option(
"-b",
"--before",
default=MIN_DATETIME,
help="License is invalid before this date.",
type=click.DateTime(formats=DATETIME_FORMATS),
)
@click.option(
"-a",
"--after",
default=MAX_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"""
Expand All @@ -251,5 +328,5 @@ def generate_license(ctx, key, mac, ipv4, before, after):
click.echo(FINISH_GENERATE_LICENSE_MSG)


if __name__ == '__main__':
if __name__ == "__main__":
cli()
12 changes: 7 additions & 5 deletions pyencrypt/decrypt.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,22 +6,24 @@

def decrypt_key(cipher_key: str, d: int, n: int) -> str:
plain_ls = list()
for num in map(int, cipher_key.split('O')):
for num in map(int, cipher_key.split("O")):
plain_ls.append(pow(num, d, n))
# 去掉intt后末尾多余的0
return ''.join(map(chr, filter(lambda x: x != 0, intt(plain_ls))))
return "".join(map(chr, filter(lambda x: x != 0, intt(plain_ls))))


def _decrypt_file(data: bytes, key: str) -> bytes:
return aes_decrypt(data, key)


def decrypt_file(path: Path, key: str, delete_origin: bool = False, new_path: Path = None) -> bytes:
if path.suffix != '.pye':
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)
if new_path:
if new_path.suffix != '.py':
if new_path.suffix != ".py":
raise Exception("Origin file path must be py suffix.")
new_path.touch(exist_ok=True)
new_path.write_bytes(data)
Expand Down
Loading

0 comments on commit 4a0bc4b

Please sign in to comment.