99 lines
2.9 KiB
Python
99 lines
2.9 KiB
Python
import io
|
|
import os
|
|
import struct
|
|
import tempfile
|
|
|
|
from Crypto.Cipher import AES
|
|
|
|
|
|
def encrypt_file(key, in_file, out_file, 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 or tempfile
|
|
|
|
out_file:
|
|
Name of the output file or tempfile
|
|
|
|
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 type(in_file) is str:
|
|
filesize = os.path.getsize(in_file)
|
|
input = 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)
|
|
input = in_file
|
|
else:
|
|
print("Unable to encrypt Input " + str(in_file) + " of type " + str(type(in_file)))
|
|
return
|
|
|
|
if type(out_file) is str:
|
|
output = open(out_file, 'wb')
|
|
|
|
elif type(out_file) is io.BufferedRandom or tempfile.SpooledTemporaryFile:
|
|
output = out_file
|
|
else:
|
|
print("Unable to encrypt to Output " + str(out_file) + " of type " + str(type(out_file)))
|
|
return
|
|
|
|
cipher = AES.new(key, AES.MODE_CBC)
|
|
|
|
output.write(struct.pack('<Q', filesize))
|
|
output.write(cipher.iv)
|
|
|
|
while chunk := input.read(chunksize):
|
|
if len(chunk) % 16 != 0:
|
|
chunk += b' ' * (16 - len(chunk) % 16)
|
|
output.write(cipher.encrypt(chunk))
|
|
|
|
|
|
def decrypt_file(key, in_file, out_file, 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 type(in_file) is str:
|
|
input = open(in_file, 'rb')
|
|
|
|
elif type(in_file) is tempfile.SpooledTemporaryFile:
|
|
in_file.seek(0)
|
|
input = in_file
|
|
else:
|
|
print("Unable to unencrypt Input " + str(in_file) + " of type " + str(type(in_file)))
|
|
return
|
|
|
|
if type(out_file) is str:
|
|
output = open(out_file, 'wb')
|
|
|
|
elif type(out_file) is io.BufferedRandom or tempfile.SpooledTemporaryFile:
|
|
output = out_file
|
|
else:
|
|
print("Unable to unencrypt to Output " + str(out_file) + " of type " + str(type(out_file)))
|
|
return
|
|
|
|
origsize = struct.unpack('<Q', input.read(struct.calcsize('Q')))[0]
|
|
iv = input.read(16)
|
|
decryptor = AES.new(key, AES.MODE_CBC, iv)
|
|
|
|
while chunk := input.read(chunksize):
|
|
output.write(decryptor.decrypt(chunk))
|
|
|
|
output.truncate(origsize)
|