import os
import random
import sys
# 随机生成密匙random.randint(1, 255)
_key = random.randint(1, 255)
# 加密过的标示,两个字节
have_encrypt_first_flag = 0x32
have_encrypt_second_flag = 0x65
have_encrypt_third_flag = 0x28
# images_path是需要加密的图片的绝对路径
def encrypt(image_path, en_key=_key):
    f = open(image_path, 'rb')
    file_data = f.read()
    file_size = f.tell()
    print("file size====", file_size)
    f.close()
    file_byte_array = bytearray(file_data)
    # 判断是否加密过
    if (have_encrypt_first_flag == file_byte_array[0]
        and have_encrypt_second_flag == file_byte_array[1]
        and have_encrypt_third_flag == file_byte_array[2]):
        print('this picture has encrypt before!')
    else:
        print("开始加密文件=============", image_path)
        encrypt_file_byte_array = bytearray(0)
        encrypt_file_byte_array.append(have_encrypt_first_flag)
        encrypt_file_byte_array.append(have_encrypt_second_flag)
        encrypt_file_byte_array.append(have_encrypt_third_flag)
        encrypt_file_byte_array.append(en_key)
        for byte in file_byte_array:
            encrypt_byte = byte ^ en_key
            encrypt_file_byte_array.append(encrypt_byte)
        os.remove(image_path)
        f2 = open(image_path, 'wb')
        f2.write(encrypt_file_byte_array)
        f2.close()
        print("加密文件结束=============", image_path)
def decrypt(image_path, en_key=_key):
    f = open(image_path, 'rb')
    file_data = f.read()
    file_size = f.tell()
    print("file size====", file_size)
    f.close()
    file_byte_array = bytearray(file_data)
    # 判断是否加密过
    if (have_encrypt_first_flag == file_byte_array[0] and have_encrypt_second_flag == file_byte_array[1]
        and have_encrypt_third_flag == file_byte_array[2]):
        print("开始解密文件=============", image_path)
        file_byte_array.remove(have_encrypt_first_flag)
        file_byte_array.remove(have_encrypt_second_flag)
        file_byte_array.remove(have_encrypt_third_flag)
        file_byte_array.remove(en_key)
        encrypt_file_byte_array = bytearray(0)
        for byte in file_byte_array:
            encrypt_byte = byte ^ en_key
            encrypt_file_byte_array.append(encrypt_byte)
        os.remove(image_path)
        f2 = open(image_path, 'wb')
        f2.write(encrypt_file_byte_array)
        f2.close()
        print("解密文件结束=============", image_path)
    else:
        print('this picture has decrypt before!')
# 判断是否是png和jpg图片
def judge_is_png_or_jpg(file_name, suffix_list=['.jpg', '.png']):
    suffix = os.path.splitext(file_name)[1].lower()
    if suffix in suffix_list:
        return True
    else:
        return False
def get_file_path_list(file_dir, suffix_list=['.jpg', '.png']):
    file_path_list = []
    for root, dirs, files in os.walk(file_dir):
        for file in files:
            file_suffix = os.path.splitext(file)[1]
            if file_suffix in suffix_list:
                file_path_list.append(os.path.join(root, file))
    return file_path_list
def encrypt_img(image_path, encrypt_flag, encrypt_key):
    if os.path.isfile(image_path):
        if not judge_is_png_or_jpg(image_path):
            print('请确认文件类型是否为图片,目前只支持png和jpg类型')
        # 解密
        if encrypt_flag:
            encrypt(image_path, encrypt_key)
        # 加密
        else:
            decrypt(image_path, encrypt_key)
    elif os.path.isdir(image_path):
        path_list = get_file_path_list(image_path)
        for file_path in path_list:
            # 解密
            if encrypt_flag:
                encrypt(file_path, encrypt_key)
            # 加密
            else:
                decrypt(file_path, encrypt_key)
    else:
        print("无效路径!请确认该路径是否存在[{}]".format(image_path))
if __name__ == '__main__':
    """python encrypt_img.py xxx -e/-d xxx """
    args = sys.argv
    if args.__len__() > 1:
        if args[1] == 'help':
            print("""
                    function: encrypt_img
                    params:
                        1: 文件路径,或者目录
                        2: 加密或者解密
                        3: 传入密钥(可不传)
                    """)
        else:
            _path = args[1]
            # 解密
            if args.__len__() > 2 and args[2] == '-d':
                _flag = False
            # 加密
            else:
                _flag = True
            if args.__len__() > 3:
                if str(args[3]).isdigit():
                    _key = int(args[3])
                else:
                    print("加密密钥必须为整数!")
            encrypt_img(_path, _flag, _key)
    else:
        print("\n请传入必要参数!如需帮助请输入\nencrypt_img.py help")