Python工具箱系列(十三)

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

我們如下約定：

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

使用python經典的命令行框架argparse，它是python標準庫的一部分，也就是說安裝好python3.8就自帶這個框架，不需要再安裝什麼。程式員經常說：talk is cheap,show me your code。所以將代碼如下所示。

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

通過命令行參數的使用，可以為用戶提供了最大的靈活性。通過命令行，再結合SHELL腳本等，可以批量處理大量的文件、目錄，因此命令行仍然是專家、高手的最愛。