1 files changed, 71 insertions, 0 deletions
diff --git a/filecrypt.py b/filecrypt.py
new file mode 100644
index 0000000..ce6e3a1
--- /dev/null
+++ b/filecrypt.py
@@ -0,0 +1,71 @@
+import secrets
+from Cryptodome.Cipher import AES
+import base64
+import struct
+import subprocess
+import binascii
+import os
+
+FILE_ID_LENGTH = 22
+TOKEN_LENGTH = 22
+
+def encrypt_file(filename_in, chunk_size=1000000//16):
+    file_id = secrets.token_urlsafe(16)
+    auth_secret = secrets.token_bytes(16)
+    key = secrets.token_bytes(16)
+    data_nonce = secrets.token_bytes(8)
+
+    token_cipher = AES.new(auth_secret, AES.MODE_GCM)
+    ciphertext, token_tag = token_cipher.encrypt_and_digest(key)
+    token = base64.b64encode(ciphertext).rstrip(b'=').decode()
+
+    with open(f'{file_id}.enc', 'wb') as fout, open(filename_in, 'rb') as fin:
+        fout.write(token_cipher.nonce) # 16 bytes
+        fout.write(token_tag) # 16 bytes
+        fout.write(auth_secret) # 16 bytes
+        fout.write(data_nonce) # 8 bytes
+
+        cipher = AES.new(key, AES.MODE_CTR,
+                initial_value=struct.pack('<Q', 0), nonce=data_nonce)
+
+        block = fin.read(cipher.block_size*chunk_size)
+        while block:
+            data = cipher.encrypt(block)
+            fout.write(data)
+            block = fin.read(cipher.block_size*chunk_size)
+
+    return file_id, token
+
+def payload_size(path):
+    return os.stat(path).st_size - 56
+
+def decrypt_generator(filename, token, seek=0, end=None, chunk_size=1000000//16):
+    with open(filename, 'rb') as fin:
+        token_nonce = fin.read(16)
+        token_tag = fin.read(16)
+        auth_secret = fin.read(16)
+        data_nonce = fin.read(8)
+
+        ciphertext = base64.b64decode(token + '='*((3-len(token)%3)%3))
+        token_cipher = AES.new(auth_secret, AES.MODE_GCM, nonce=token_nonce)
+        key = token_cipher.decrypt_and_verify(ciphertext, token_tag)
+
+        # Use token cipher's block size assuming the two are the same
+        cipher = AES.new(key, AES.MODE_CTR,
+                initial_value=struct.pack('>Q', seek // token_cipher.block_size), nonce=data_nonce)
+        print('Seeking to position', seek, 'iv', seek // token_cipher.block_size, 'pos', seek - seek %
+                cipher.block_size, 'offset', seek % cipher.block_size)
+        fin.seek(seek - seek % cipher.block_size, 1)
+        offset = seek % cipher.block_size
+
+        to_send = end - seek + 1 if end else None
+        block = fin.read(cipher.block_size*chunk_size)
+        while block and (to_send is None or to_send > 0):
+            data = cipher.decrypt(block)
+            out = data[offset:to_send]
+            yield out
+            if to_send is not None:
+                to_send -= len(out)
+            offset = 0
+            block = fin.read(cipher.block_size*chunk_size)
+