Compress and Crypt

compress.py

@@ -1,6 +1,8 @@
 import os
 import re
 import gzip
+import tempfile
+from crypt import *
 
 
 def recurse(path, maxsize="50M"):
@@ -12,17 +14,20 @@ def recurse(path, maxsize="50M"):
             Hsize = str(human_readable_size(size, 0))
             print(f + " : " + Hsize)
             if size > max:
-                compress(uri)
+                compress(uri, key)
         elif os.path.isdir(uri):
             recurse(uri)
 
 
-def compress(file):
+def compress(file, key):
+    gz = tempfile.SpooledTemporaryFile()
     with open(file, 'rb') as infile:
-        with gzip.open(file + ".gz", 'wb') as zipfile:
+        with gzip.open(gz, 'wb') as zipfile:
             while chunk := infile.read(64 * 1024):
                 zipfile.write(chunk)
-    print(file + " compressed")
+        print(file + " compressed")
+        encrypt_file(key, gz, file + ".gz.enc")
+        print(file + " encrypted")
     pass
 
 

crypt.py

@@ -1,6 +1,11 @@
-import os, random, struct
+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.
@@ -22,22 +27,35 @@ def encrypt_file(key, in_file, out_file=None, chunksize=64 * 1024):
             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)
-    filesize = os.path.getsize(in_file)
 
-    with open(in_file, 'rb') as infile:
-        with open(out_file, 'wb') as outfile:
-            outfile.write(struct.pack('<Q', filesize))
-            outfile.write(cipher.iv)
+    if type(in_file) is str:
+        filesize = os.path.getsize(in_file)
+        infile = open(in_file, 'rb')
 
-            while chunk := infile.read(chunksize):
-                if len(chunk) % 16 != 0:
-                    chunk += b' ' * (16 - len(chunk) % 16)
+    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))
 
-                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

main.py

@@ -1,17 +1,24 @@
-from crypt import *
 from Crypto.Random import get_random_bytes
-
+from crypt import *
 from compress import *
 
-key = get_random_bytes(32) # 32 bytes * 8 = 256 bits (1 byte = 8 bits)
-print(key)
+if not os.path.exists("key"):
+    key = get_random_bytes(32)  # 32 bytes * 8 = 256 bits (1 byte = 8 bits)
+    open("key", 'wb').write(key)
+    print("Saved")
+else:
+    key = open("key", 'rb').read()
+    print("Recovered")
 
+print(key)
 file = "original_file.txt"
-cfile = file+'.enc'
+cfile = file + '.enc'
 
 encrypt_file(key, file, cfile)
 
 decrypt_file(key, cfile)
 
-rootdir="test"
-recurse(rootdir)
+rootdir = "test"
+recurse(rootdir, key)
+
+decrypt_file(key, "test/DressLens4.mp4.gz.enc", "test/DressLens4.mp4.gz")