import argparse
import os
import struct
import sys
from pathlib import Path

from Crypto.Cipher import AES
from Crypto.Util.Padding import pad, unpad

__authors__ = 'tianbin'
__version__ = 'version 0.9'
__license__ = 'free'

defaultsize = 64*1024


def encrypt_file(key, in_filename, out_filename=None, chunksize=defaultsize):
    """
    对文件进行加密

    Args:
        key (str): 16字节密钥
        in_filename (str): 待加密文件
        out_filename (str, optional): 加密后输出的文件
        chunksize (int, optional): 块大小,缺省64k
    """
    if not out_filename:
        out_filename = in_filename + '.enc'
    iv = os.urandom(16)
    encryptor = AES.new(key, AES.MODE_CBC, iv)
    filesize = os.path.getsize(in_filename)
    with open(in_filename, 'rb') as infile:
        with open(out_filename, 'wb') as outfile:
            outfile.write(struct.pack('<Q', filesize))
            outfile.write(iv)
            pos = 0
            while pos < filesize:
                chunk = infile.read(chunksize)
                pos += len(chunk)
                if pos == filesize:
                    chunk = pad(chunk, AES.block_size)
                outfile.write(encryptor.encrypt(chunk))


def decrypt_file(key, in_filename, out_filename=None, chunksize=defaultsize):
    """
    解密文件

    Args:
        key (str): 16字节密钥
        in_filename (str): 待解密文件
        out_filename (str, optional): 解密后输出的文件
        chunksize (int, optional): 块大小,缺省64K
    """
    if not out_filename:
        out_filename = in_filename + '.dec'
    with open(in_filename, 'rb') as infile:
        filesize = struct.unpack('<Q', infile.read(8))[0]
        iv = infile.read(16)
        encryptor = AES.new(key, AES.MODE_CBC, iv)
        with open(out_filename, 'wb') as outfile:
            encrypted_filesize = os.path.getsize(in_filename)
            pos = 8 + 16  # the filesize and IV.
            while pos < encrypted_filesize:
                chunk = infile.read(chunksize)
                pos += len(chunk)
                chunk = encryptor.decrypt(chunk)
                if pos == encrypted_filesize:
                    chunk = unpad(chunk, AES.block_size)
                outfile.write(chunk)


if __name__ == '__main__':
    
    # 密钥随便写，使用时只使用前16字节
    key = 'stayhungrystayfoolish'
    realkey = key[:16].encode('utf-8')

    def parser():
        """
        分析用户命令行
        """
        parser = argparse.ArgumentParser()
        parser.add_argument("-d", "--decry", action="store_true",
                            help="解密模式")
        parser.add_argument("-e", "--encry", action="store_true",
                            help="加密模式")
        parser.add_argument("-i", "--input", type=str,
                            help="要处理的文件")
        parser.add_argument("-o", "--output", type=str,
                            help="要输出的文件")

        args = parser.parse_args()
        print(args)
        # 判断参数输入情况,如果没有参数,则显示帮助。
        if len(sys.argv) == 1:
            parser.print_help()
            return

        # 解密模式，获得输入与输出文件后，调用算法解密
        if args.decry:
            inputfilename = Path(args.input)
            if inputfilename.exists():
                decrypt_file(realkey,in_filename=args.input,out_filename=args.output)
            else:
                print(f'{args.input}不存在')

        # 加密模式，获得输入与输出文件后，调用算法加密
        if args.encry:
            inputfilename = Path(args.input)
            if inputfilename.exists():
                encrypt_file(realkey,in_filename=args.input,out_filename=args.output)
            else:
                print(f'{args.input}不存在')

    parser()

# 以下命令显示帮助信息
python .\aeshandler.py                         
usage: aeshandler.py [-h] [-d] [-e] [-i INPUT] [-o OUTPUT]

optional arguments:
  -h, --help            show this help message and exit
  -d, --decry           解密模式
  -e, --encry           加密模式
  -i INPUT, --input INPUT
                        要处理的文件
  -o OUTPUT, --output OUTPUT
                        要输出的文件

# 以下命令加密指定的文件，加密后的文件为test1.docx
python .\aeshandler.py -e -i ../resources/神龟虽寿.docx -o test1.docx
Namespace(decry=False, encry=True, input='../resources/神龟虽寿.docx', output='test1.docx')

# 以下命令解密指定的文件,要解密的文件为test1.docx，解密后的文件为test2.docx
python .\aeshandler.py -d -i test1.docx -o test2.docx                
Namespace(decry=True, encry=False, input='test1.docx', output='test2.docx')