From 5c2d2a22a6abea84952a755c36724c811799ebb5 Mon Sep 17 00:00:00 2001
From: LPOCHOLLE <>
Date: Mon, 27 Sep 2021 15:56:05 +0200
Subject: [PATCH] Implement Backup class

---
 Backup.py                                   | 91 +++++++++++++++++++++
 compress.py                                 | 47 -----------
 crypt.py                                    | 78 +++++++++++-------
 main.py                                     | 16 ++--
 original_file.txt => test/original_file.txt |  0
 5 files changed, 143 insertions(+), 89 deletions(-)
 create mode 100644 Backup.py
 delete mode 100644 compress.py
 rename original_file.txt => test/original_file.txt (100%)

diff --git a/Backup.py b/Backup.py
new file mode 100644
index 0000000..5dadbee
--- /dev/null
+++ b/Backup.py
@@ -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])
diff --git a/compress.py b/compress.py
deleted file mode 100644
index a72deca..0000000
--- a/compress.py
+++ /dev/null
@@ -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])
\ No newline at end of file
diff --git a/crypt.py b/crypt.py
index 2a9c608..d5d24b8 100644
--- a/crypt.py
+++ b/crypt.py
@@ -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):
-            if len(chunk) % 16 != 0:
-                chunk += b' ' * (16 - len(chunk) % 16)
+    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
 
-            outfile.write(cipher.encrypt(chunk))
+    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=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)
-        decryptor = AES.new(key, AES.MODE_CBC, iv)
+    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
 
-        with open(out_file, 'wb') as outfile:
-            while chunk := infile.read(chunksize):
-                outfile.write(decryptor.decrypt(chunk))
+    if type(out_file) is str:
+        output = open(out_file, 'wb')
 
-            outfile.truncate(origsize)
\ No newline at end of file
+    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)
diff --git a/main.py b/main.py
index 47d45e5..7f10151 100644
--- a/main.py
+++ b/main.py
@@ -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")
\ No newline at end of file
+bck.recurse(rootdir)
diff --git a/original_file.txt b/test/original_file.txt
similarity index 100%
rename from original_file.txt
rename to test/original_file.txt