crypto.py

import os 
import hashlib
import time
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
from datetime import datetime

SALT_SIZE = 16 
IV_SIZE = 12
ITERATIONS = 600000

def derive_key(password, salt):
    kdf = PBKDF2HMAC(
        algorithm=hashes.SHA256(),
        length = 32,
        salt=salt,
        iterations=ITERATIONS,
    )
    return kdf.derive(password.encode())

def encrypt_file(input_path, output_path, password):
    start = time.time()

    salt = os.urandom(SALT_SIZE)
    iv = os.urandom(IV_SIZE)
    key = derive_key(password, salt) 

    with open(input_path, "rb") as f:
        plaintext = f.read()

    aesgcm = AESGCM(key)
    ciphertext = aesgcm.encrypt(iv, plaintext, None)

    with open(output_path, "wb") as f:
        f.write(salt)
        f.write(iv)
        f.write(ciphertext)

    return round(time.time() - start, 4) #return seconds

def decrypt_file(input_path, output_path, password):
    start = time.time()

    with open(input_path, "rb") as f:
        data = f.read()
    
    salt = data[:SALT_SIZE]
    iv   = data[SALT_SIZE:SALT_SIZE + IV_SIZE]
    ciphertext = data[SALT_SIZE + IV_SIZE:]

    key = derive_key(password, salt)

    try:
        aesgcm = AESGCM(key)
        plaintext = aesgcm.decrypt(iv, ciphertext, None)
    except Exception: 
        raise ValueError("Wrong password or file is corrupted.")
    
    with open(output_path, "wb") as f:
        f.write(plaintext)

    return round (time.time() - start, 4) #return seconds 

def get_file_hash(filepath):
    h = hashlib.sha256()
    with open(filepath, "rb") as f:
        while chunk := f.read(8192):
            h.update(chunk)
        return h.hexdigest()
    
def write_log(action, filepath, duration):
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    filename = os.path.basename(filepath) 
    with open("encryption_log.txt", "a") as f:
        f.write(f"[{timestamp}] {action} | {filename} | {duration}s\n")


if __name__ == "__main__": 
    encrypt_file("test.txt", "test.enc", "mypassword")
    print("Encrypted OK")

    decrypt_file("test.enc", "test_decrypted.txt", "mypassword")
    print("Decrypted OK")