Implement Backup class

This commit is contained in:
LPOCHOLLE 2021-09-27 15:56:05 +02:00
parent 274711a22d
commit 5c2d2a22a6
5 changed files with 143 additions and 89 deletions

91
Backup.py Normal file
View File

@ -0,0 +1,91 @@
import re
import gzip
from crypt import *
class Backup:
key = None
compressMinSize = "50M"
save_location = None
def __init__(self, key):
self.key = key
def recurse(self, path):
min_size = parse_size(self.compressMinSize)
for f in os.listdir(path):
uri = os.path.join(path, f)
if os.path.isfile(uri):
size = os.path.getsize(uri)
print(f + " : ", human_size(size))
if size > min_size:
enc = crypt(compress(uri), self.key)
else:
enc = crypt(uri, self.key)
save(enc, os.path.join(self.save_location, f + ".enc"))
elif os.path.isdir(uri):
self.recurse(uri)
def compress(file):
if type(file) is str:
infile = open(file, 'rb')
elif type(file) is io.BufferedRandom or tempfile.SpooledTemporaryFile:
file.seek(0)
infile = file
compressed_file = tempfile.SpooledTemporaryFile()
with gzip.open(compressed_file, 'wb') as zipfile:
while chunk := infile.read(64 * 1024):
zipfile.write(chunk)
return compressed_file
def uncompress(file):
if type(file) is io.BufferedRandom or tempfile.SpooledTemporaryFile:
file.seek(0)
decompressed_file = tempfile.SpooledTemporaryFile()
with gzip.open(file, 'rb') as zipfile:
while chunk := zipfile.read(64 * 1024):
decompressed_file.write(chunk)
return decompressed_file
def crypt(file, key):
encrypted_file = tempfile.SpooledTemporaryFile()
encrypt_file(key, file, encrypted_file)
return encrypted_file
def uncrypt(file, key):
decrypted_file = tempfile.SpooledTemporaryFile()
decrypt_file(key, file, decrypted_file)
return decrypted_file
def save(file, name):
if type(file) is io.BufferedRandom or tempfile.SpooledTemporaryFile:
file.seek(0)
with open(name, 'wb') as save:
while chunk := file.read(64 * 1024):
save.write(chunk)
else:
print("Unable to save " + str(file) + " of type " + str(type(file)))
return
def human_size(size, decimal_places=0):
for unit in ['B', 'K', 'M', 'G', 'T']:
if size < 1024.0:
break
size /= 1024.0
return f"{size:.{decimal_places}f}{unit}"
def parse_size(size):
units = {"B": 1, "K": 2**10, "M": 2**20, "G": 2**30, "T": 2**40}
if size[-1].isdigit():
size = size + 'K'
number, unit = re.match(r"([0-9]+)([BKMGT])", size, re.I).groups()
return int(float(number)*units[unit])

View File

@ -1,47 +0,0 @@
import os
import re
import gzip
import tempfile
from crypt import *
def recurse(path, maxsize="50M"):
max = parse_size(maxsize)
for f in os.listdir(path):
uri = os.path.join(path, f)
if os.path.isfile(uri):
size = os.path.getsize(uri)
Hsize = str(human_readable_size(size, 0))
print(f + " : " + Hsize)
if size > max:
compress(uri, key)
elif os.path.isdir(uri):
recurse(uri)
def compress(file, key):
gz = tempfile.SpooledTemporaryFile()
with open(file, 'rb') as infile:
with gzip.open(gz, 'wb') as zipfile:
while chunk := infile.read(64 * 1024):
zipfile.write(chunk)
print(file + " compressed")
encrypt_file(key, gz, file + ".gz.enc")
print(file + " encrypted")
pass
def human_readable_size(size, decimal_places=2):
for unit in ['B', 'K', 'M', 'G', 'T']:
if size < 1024.0:
break
size /= 1024.0
return f"{size:.{decimal_places}f}{unit}"
def parse_size(size):
units = {"B": 1, "K": 2**10, "M": 2**20, "G": 2**30, "T": 2**40}
if size[-1].isdigit():
size = size + 'K'
number, unit = re.match(r"([0-9]+)([BKMGT])", size, re.I).groups()
return int(float(number)*units[unit])

View File

@ -6,7 +6,7 @@ import tempfile
from Crypto.Cipher import AES
def encrypt_file(key, in_file, out_file=None, chunksize=64 * 1024):
def encrypt_file(key, in_file, out_file, chunksize=64 * 1024):
""" Encrypts a file using AES (CBC mode) with the
given key.
@ -16,10 +16,10 @@ def encrypt_file(key, in_file, out_file=None, chunksize=64 * 1024):
are more secure.
in_file:
Name of the input file
Name of the input file or tempfile
out_file:
If None, '<in_filename>.enc' will be used.
Name of the output file or tempfile
chunksize:
Sets the size of the chunk which the function
@ -28,36 +28,40 @@ def encrypt_file(key, in_file, out_file=None, chunksize=64 * 1024):
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')
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)
infile = in_file
input = in_file
else:
print("Unable to encrypt " + str(in_file) + " of type " + str(type(in_file)))
print("Unable to encrypt Input " + 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)
if type(out_file) is str:
output = open(out_file, 'wb')
while chunk := infile.read(chunksize):
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)
outfile.write(cipher.encrypt(chunk))
output.write(cipher.encrypt(chunk))
def decrypt_file(key, in_file, out_file=None, chunksize=24 * 1024):
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
@ -65,18 +69,30 @@ def decrypt_file(key, in_file, out_file=None, chunksize=24 * 1024):
(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'
if type(in_file) is str:
input = open(in_file, 'rb')
with open(in_file, 'rb') as infile:
origsize = struct.unpack('<Q', infile.read(struct.calcsize('Q')))[0]
iv = infile.read(16)
elif type(in_file) is io.BufferedRandom or 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)
with open(out_file, 'wb') as outfile:
while chunk := infile.read(chunksize):
outfile.write(decryptor.decrypt(chunk))
while chunk := input.read(chunksize):
output.write(decryptor.decrypt(chunk))
outfile.truncate(origsize)
output.truncate(origsize)

16
main.py
View File

@ -1,6 +1,6 @@
from Crypto.Random import get_random_bytes
from crypt import *
from compress import *
import os
from Backup import Backup
if not os.path.exists("key"):
key = get_random_bytes(32) # 32 bytes * 8 = 256 bits (1 byte = 8 bits)
@ -11,14 +11,8 @@ else:
print("Recovered")
print(key)
file = "original_file.txt"
cfile = file + '.enc'
encrypt_file(key, file, cfile)
decrypt_file(key, cfile)
bck = Backup(key)
bck.save_location = "C:\\temp"
rootdir = "test"
recurse(rootdir, key)
decrypt_file(key, "test/DressLens4.mp4.gz.enc", "test/DressLens4.mp4.gz")
bck.recurse(rootdir)