82 lines
2.4 KiB
Python
82 lines
2.4 KiB
Python
import io
|
|
import os
|
|
import struct
|
|
import tempfile
|
|
|
|
from Crypto.Cipher import AES
|
|
|
|
|
|
def encrypt_file(key, in_file, out_file=None, chunksize=64 * 1024):
|
|
""" Encrypts a file using AES (CBC mode) with the
|
|
given key.
|
|
|
|
key:
|
|
The encryption key - a string that must be
|
|
either 16, 24 or 32 bytes long. Longer keys
|
|
are more secure.
|
|
|
|
in_file:
|
|
Name of the input file
|
|
|
|
out_file:
|
|
If None, '<in_filename>.enc' will be used.
|
|
|
|
chunksize:
|
|
Sets the size of the chunk which the function
|
|
uses to read and encrypt the file. Larger chunk
|
|
sizes can be faster for some files and machines.
|
|
chunksize must be divisible by 16.
|
|
"""
|
|
|
|
if not out_file:
|
|
out_file = in_file + '.enc'
|
|
|
|
cipher = AES.new(key, AES.MODE_CBC)
|
|
|
|
if type(in_file) is str:
|
|
filesize = os.path.getsize(in_file)
|
|
infile = open(in_file, 'rb')
|
|
|
|
elif type(in_file) is io.BufferedRandom or tempfile.SpooledTemporaryFile:
|
|
in_file.seek(0, 2)
|
|
filesize = in_file.tell()
|
|
in_file.seek(0)
|
|
infile = in_file
|
|
else:
|
|
print("Unable to encrypt " + str(in_file) + " of type " + str(type(in_file)))
|
|
return
|
|
|
|
with open(out_file, 'wb') as outfile:
|
|
outfile.write(struct.pack('<Q', filesize))
|
|
outfile.write(cipher.iv)
|
|
|
|
while chunk := infile.read(chunksize):
|
|
if len(chunk) % 16 != 0:
|
|
chunk += b' ' * (16 - len(chunk) % 16)
|
|
|
|
outfile.write(cipher.encrypt(chunk))
|
|
|
|
|
|
def decrypt_file(key, in_file, out_file=None, chunksize=24 * 1024):
|
|
""" Decrypts a file using AES (CBC mode) with the
|
|
given key. Parameters are similar to encrypt_file,
|
|
with one difference: out_filename, if not supplied
|
|
will be in_filename without its last extension
|
|
(i.e. if in_filename is 'aaa.zip.enc' then
|
|
out_filename will be 'aaa.zip')
|
|
"""
|
|
#if not out_file:
|
|
# out_file = os.path.splitext(in_file)[0]
|
|
if not out_file:
|
|
out_file = in_file + '.dec'
|
|
|
|
with open(in_file, 'rb') as infile:
|
|
origsize = struct.unpack('<Q', infile.read(struct.calcsize('Q')))[0]
|
|
iv = infile.read(16)
|
|
decryptor = AES.new(key, AES.MODE_CBC, iv)
|
|
|
|
with open(out_file, 'wb') as outfile:
|
|
while chunk := infile.read(chunksize):
|
|
outfile.write(decryptor.decrypt(chunk))
|
|
|
|
outfile.truncate(origsize) |