Python工具箱系列(十二)

import os
import struct
from Crypto.Cipher import AES
from Crypto.Util.Padding import pad, unpad
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)
# 密钥，随便写，使用时只使用前16字节
key = 'stayhungrystayfoolish'
def test_enc():
    sourcefilename=r'../resources/神龟虽寿.docx'
    encrypt_file(key[:16].encode('utf-8'), sourcefilename)
def test_dec():
    sourcefilename=r'../resources/神龟虽寿.docx.enc'
    targetfilename=r'../resources/decrypt.docx'
    decrypt_file(key[:16].encode('utf-8'), sourcefilename,out_filename=targetfilename)
if __name__ == '__main__':
    test_enc()
    test_dec()

pip install PyCryptodome