Python工具箱系列(十三)

上文介紹了使用AES算法進行檔案加解密的代碼。但是如果在代碼中寫死了(hardcode)檔案名，每次要加解密檔案都要去改python源代碼，顯然有些太笨了。為此，可以使用指令行參數來在不改動源代碼的情況下，對指令行參數所指定的檔案進行加/解密操作。也可以指定加解密後輸出的檔案名稱，以友善使用。

我們如下約定：

• python檔案名為aeshandler.py
• -i,表示輸入檔案名
• -o,表示輸出檔案名
• -e,表示加密
• -d,表示解密

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')