From 62272039d6354b494217229fc8516248169dd42b Mon Sep 17 00:00:00 2001 From: Tiago Garcia Date: Wed, 11 Dec 2024 22:57:41 +0000 Subject: [PATCH] Fix file encryption Signed-off-by: Tiago Garcia --- delivery2/client/bin/rep_add_doc | 4 +- delivery2/lib/symmetric_encryption.py | 58 +++++++++++++++------------ delivery2/server/models/file.py | 2 - delivery2/server/routes/file.py | 4 +- delivery2/server/services/files.py | 3 +- 5 files changed, 38 insertions(+), 33 deletions(-) diff --git a/delivery2/client/bin/rep_add_doc b/delivery2/client/bin/rep_add_doc index 8576080..2c5d964 100755 --- a/delivery2/client/bin/rep_add_doc +++ b/delivery2/client/bin/rep_add_doc @@ -51,10 +51,10 @@ def addDoc(args): args.session = json.load(f) #Encrypt content - key, content, nonce = encrypt_file(BASE_DIR + args.file, BASE_DIR + 'encryptedText') + key, content = encrypt_file(BASE_DIR + args.file, BASE_DIR + 'encryptedText') #Upload document metadata - doc = {'document_name' : args.name, 'key' : key.hex(), 'alg' : 'AES-CFB', 'nonce' : nonce.hex() } + doc = {'document_name' : args.name, 'key' : key.hex(), 'alg' : 'AES-CFB' } try: req = requests.post(f'http://{state['REP_ADDRESS']}/file/upload/metadata', json=json.dumps(doc), diff --git a/delivery2/lib/symmetric_encryption.py b/delivery2/lib/symmetric_encryption.py index 094e6a3..040a697 100644 --- a/delivery2/lib/symmetric_encryption.py +++ b/delivery2/lib/symmetric_encryption.py @@ -6,41 +6,49 @@ from cryptography.hazmat.backends import default_backend # Function to encrypt a file using a salt def encrypt_file(input_file, output_file=None): - key = os.urandom(16) - iv = os.urandom(16) + key = os.urandom(16) + iv = os.urandom(16) - cipher = Cipher(algorithms.AES(key), modes.CFB(iv)) - encryptor = cipher.encryptor() + cipher = Cipher(algorithms.AES(key), modes.CFB(iv)) + encryptor = cipher.encryptor() - with open(input_file, 'rb') as f: - plaintext = f.read() + encrypted_content = b"" - ciphertext = encryptor.update(plaintext) + encryptor.finalize() - ciphertext = iv + ciphertext + if output_file is not None: + with open(input_file, 'rb') as infile, open(output_file, 'wb') as outfile: + # Write the IV to the output file first + outfile.write(iv) + encrypted_content += iv - if output_file is not None: - with open(output_file, 'wb') as f: - f.write(ciphertext) + while chunk := infile.read(2048): + ciphertext = encryptor.update(chunk) + outfile.write(ciphertext) + encrypted_content += ciphertext - print(iv.hex()) + # Finalize encryption + final_chunk = encryptor.finalize() + outfile.write(final_chunk) + encrypted_content += final_chunk - return key, ciphertext, iv + return key, encrypted_content # Function to decrypt a file -def decrypt_file(nonce, key, input_file, output_file=None): - with open(input_file, 'rb') as f: - encrypted_data = f.read() +def decrypt_file(key, input_file, output_file=None): + with open(input_file, 'rb') as infile: + # Read the IV from the input file + iv = infile.read(16) + cipher = Cipher(algorithms.AES(key), modes.CFB(iv)) + decryptor = cipher.decryptor() - ciphertext = encrypted_data + if output_file is not None: + with open(output_file, 'wb') as outfile: + while chunk := infile.read(2048): + plaintext = decryptor.update(chunk) + outfile.write(plaintext) - cipher = Cipher(algorithms.AES(key), modes.CFB(nonce)) - decryptor = cipher.decryptor() + # Finalize decryption + outfile.write(decryptor.finalize()) - plaintext = decryptor.update(ciphertext) + decryptor.finalize() + return True - if output_file is not None: - with open(output_file, 'wb') as f: - f.write(plaintext) - - return plaintext.hex() diff --git a/delivery2/server/models/file.py b/delivery2/server/models/file.py index a7fc99f..8288403 100644 --- a/delivery2/server/models/file.py +++ b/delivery2/server/models/file.py @@ -11,7 +11,6 @@ class File(db_connection.Model): created_at = db_connection.Column(db_connection.Integer, nullable=False) key = db_connection.Column(db_connection.String, nullable=False) alg = db_connection.Column(db_connection.String, nullable=False) - nonce = db_connection.Column(db_connection.String, nullable=False) org_id = db_connection.Column(db_connection.Integer, db_connection.ForeignKey('organizations.id'), nullable=False) creator_id = db_connection.Column(db_connection.Integer, db_connection.ForeignKey('users.id'), nullable=False) org = db_connection.relationship('Organization', backref=db_connection.backref('org_files', uselist=False)) @@ -26,7 +25,6 @@ class File(db_connection.Model): "created_at": self.created_at, "key": self.key, "alg": self.alg, - "nonce": self.nonce, "org": {"id": self.org.id, "name": self.org.name}, "creator": {"id": self.creator.id, "username": self.creator.username} } \ No newline at end of file diff --git a/delivery2/server/routes/file.py b/delivery2/server/routes/file.py index 5eb08a6..2d5e138 100644 --- a/delivery2/server/routes/file.py +++ b/delivery2/server/routes/file.py @@ -48,7 +48,7 @@ def file_upload_metadata(): data = request.json if type(data) is str: data = json.loads(data) - if "document_name" not in data or "key" not in data or "alg" not in data or "nonce" not in data: + if "document_name" not in data or "key" not in data or "alg" not in data: return jsonify({"error": "Missing required fields"}), 400 org = OrganizationService.get_organization(session.org_id) @@ -59,7 +59,7 @@ def file_upload_metadata(): if not user: return jsonify({"error": "User not found"}), 404 - file = upload_service.create_file(session.token, org, user, data["document_name"], data["key"], data["alg"], data["nonce"]) + file = upload_service.create_file(session.token, org, user, data["document_name"], data["key"], data["alg"]) return jsonify(file.to_dict()), 201 diff --git a/delivery2/server/services/files.py b/delivery2/server/services/files.py index 9bf55f7..316409d 100644 --- a/delivery2/server/services/files.py +++ b/delivery2/server/services/files.py @@ -13,7 +13,7 @@ class FileService: def __init__(self): self.current_requests = {} - def create_file(self, session_token: str, org: Organization, user: User, file_name: str, key: str, alg: str, nonce: str) -> File: + def create_file(self, session_token: str, org: Organization, user: User, file_name: str, key: str, alg: str) -> File: file = File( file_handle = None, document_handle = get_hash(file_name), @@ -21,7 +21,6 @@ class FileService: created_at = int(datetime.now().timestamp()), key = key, alg = alg, - nonce = nonce, org_id = org.id, creator_id = user.id, org = org,