# /// script
# requires-python = ">=3.10"
# dependencies = ["cryptography"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Envelope encryption and crypto-shredding in 40 lines.

Build the envelope encryption mechanism, measure its performance, and see data deletion by key destruction, including the conflict with backups."""

import argparse
import os
import time
import json
import copy
from cryptography.hazmat.primitives.ciphers.aead import AESGCM

class KeyService:
    def __init__(self):
        self._master_key = AESGCM.generate_key(bit_length=256)
        self._aesgcm = AESGCM(self._master_key)

    def wrap(self, data_key):
        nonce = os.urandom(12)
        return nonce + self._aesgcm.encrypt(nonce, data_key, None)

    def unwrap(self, wrapped_key):
        nonce, ciphertext = wrapped_key[:12], wrapped_key[12:]
        return self._aesgcm.decrypt(nonce, ciphertext, None)

def main():
    parser = argparse.ArgumentParser(description=__doc__.split('\n')[0])
    parser.add_argument("--patients", type=int, default=5, help="Number of fake patients to generate.")
    parser.add_argument("--data-mb", type=int, default=100, help="Megabytes of data for the performance test.")
    parser.add_argument("--json", type=str, help="Also write the results to this JSON file.")
    args = parser.parse_args()

    results = {"summary": {}, "tables": {"operations": []}}
    ks = KeyService()
    db = {}
    
    print(f"--- 1. Encrypting {args.patients} Patients ---")
    for patient_id in range(1, args.patients + 1):
        data_key = AESGCM.generate_key(bit_length=256)
        wrapped_key = ks.wrap(data_key)
        
        aesgcm = AESGCM(data_key)
        nonce = os.urandom(12)
        data = f"Sensitive medical records for patient {patient_id}".encode()
        ciphertext = nonce + aesgcm.encrypt(nonce, data, None)
        
        db[patient_id] = {'wrapped_key': wrapped_key, 'ciphertext': ciphertext}
    print("Patient data encrypted and stored.")

    print(f"\n--- 2. Performance Comparison ({args.data_mb} MB) ---")
    data_key_local = AESGCM.generate_key(bit_length=256)
    aesgcm_local = AESGCM(data_key_local)
    data_bulk = os.urandom(args.data_mb * 1024 * 1024)
    nonce_bulk = os.urandom(12)
    
    t0 = time.time()
    aesgcm_local.encrypt(nonce_bulk, data_bulk, None)
    t1 = time.time()
    bulk_time = t1 - t0
    
    t2 = time.time()
    ks.wrap(data_key_local)
    t3 = time.time()
    wrap_time = t3 - t2
    
    print(f"Encrypting {args.data_mb} MB locally took: {bulk_time:.4f} seconds")
    print(f"Wrapping a 32-byte key once took: {wrap_time:.6f} seconds")
    
    results["summary"]["bulk_encrypt_time_s"] = bulk_time
    results["summary"]["key_wrap_time_s"] = wrap_time

    print("\n--- 3. Crypto-Shredding Patient 3 ---")
    # FIX: Use deepcopy so the backup actually retains its own independent wrapped_key
    backup_db = copy.deepcopy(db) 
    
    # "Erase" patient 3 by deleting only its wrapped key from the active DB
    del db[3]['wrapped_key']
    print("Patient 3's wrapped key deleted from active database.")
    
    try:
        ks.unwrap(db[3]['wrapped_key'])
        shredded = False
    except KeyError:
        print("Success: Active wrapped key is gone. Ciphertext is shredded.")
        shredded = True
        
    results["summary"]["active_db_shredded"] = shredded

    print("\n--- 4. The Backup Conflict ---")
    # The backup still has the key because of deepcopy
    try:
        recovered_key = ks.unwrap(backup_db[3]['wrapped_key'])
        backup_aesgcm = AESGCM(recovered_key)
        nonce = backup_db[3]['ciphertext'][:12]
        ciphertext = backup_db[3]['ciphertext'][12:]
        recovered_data = backup_aesgcm.decrypt(nonce, ciphertext, None)
        print(f"Data recovered from backup: {recovered_data.decode()}")
        backup_works = True
    except Exception as e:
        print(f"Failed to recover from backup: {e}")
        backup_works = False

    results["summary"]["backup_restore_works"] = backup_works
    
    results["tables"]["operations"].append({"op": "bulk_encrypt", "time_s": bulk_time})
    results["tables"]["operations"].append({"op": "key_wrap", "time_s": wrap_time})
    results["tables"]["operations"].append({"op": "shredding", "success": shredded})
    results["tables"]["operations"].append({"op": "backup_restore", "success": backup_works})

    if args.json:
        with open(args.json, "w") as f:
            json.dump(results, f, indent=2)

if __name__ == "__main__":
    main()