# /// script
# requires-python = ">=3.10"
# dependencies = []
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""A result cache that leaks, and the fix.

Reproduce the shared-cache bug and prove the fix with a deliberate verification test, including timing attacks."""

import argparse
import time
import json

DATABASE = [
    {"id": 1, "region": "north", "data": "North confidential sales: $1.2M"},
    {"id": 2, "region": "north", "data": "North confidential pipeline: $3.4M"},
    {"id": 3, "region": "south", "data": "South confidential sales: $0.8M"},
    {"id": 4, "region": "south", "data": "South confidential pipeline: $1.1M"},
]

naive_cache = {}
secure_cache = {}

def execute_query(query_text, user_region, latency):
    """Simulates a database engine with row-level security."""
    time.sleep(latency)
    if "SELECT *" in query_text:
        return [row for row in DATABASE if row["region"] == user_region]
    return []

def query_with_naive_cache(query_text, user, latency):
    """Vulnerable: Keyed only by the query text."""
    cache_key = query_text
    
    t0 = time.time()
    if cache_key in naive_cache:
        lat = time.time() - t0
        return naive_cache[cache_key], "HIT", lat
        
    result = execute_query(query_text, user["region"], latency)
    naive_cache[cache_key] = result
    lat = time.time() - t0
    return result, "MISS", lat

def query_with_secure_cache(query_text, user, latency, policy_version=1):
    """Secure: Keyed by query text, sorted roles, region, and policy version."""
    identity_tuple = (query_text, tuple(sorted(user["roles"])), user["region"], policy_version)
    cache_key = hash(identity_tuple)
    
    t0 = time.time()
    if cache_key in secure_cache:
        lat = time.time() - t0
        return secure_cache[cache_key], "HIT", lat
        
    result = execute_query(query_text, user["region"], latency)
    secure_cache[cache_key] = result
    lat = time.time() - t0
    return result, "MISS", lat

def main():
    parser = argparse.ArgumentParser(description=__doc__.split('\n')[0])
    parser.add_argument("--latency", type=float, default=0.05, help="Simulated database latency in seconds.")
    parser.add_argument("--json", type=str, help="Also write the results to this JSON file.")
    args = parser.parse_args()

    results = {"summary": {}, "tables": {"queries": []}}
    
    user_a = {"name": "User A", "region": "north", "roles": ["viewer"]}
    user_b = {"name": "User B", "region": "south", "roles": ["viewer"]}
    query = "SELECT * FROM sales"

    # --- Steps 1 & 2: The Leaking Shared Cache ---
    print("--- Steps 1 & 2: The Shared-Cache Bug ---")
    res_a, status_a, lat_a = query_with_naive_cache(query, user_a, args.latency)
    print(f"User A (North) Query -> Status: {status_a}, Rows returned: {len(res_a)}")
    results["tables"]["queries"].append({"step": "naive_user_a", "cache": "naive", "status": status_a, "latency_s": lat_a, "rows": len(res_a)})

    res_b, status_b, lat_b = query_with_naive_cache(query, user_b, args.latency)
    print(f"User B (South) Query -> Status: {status_b}, Rows returned: {len(res_b)}")
    results["tables"]["queries"].append({"step": "naive_user_b", "cache": "naive", "status": status_b, "latency_s": lat_b, "rows": len(res_b)})
    
    leak_occurred = any(r["region"] == "north" for r in res_b)
    print(f"Leak reproduced: User B received North rows? {leak_occurred}\n")
    results["summary"]["naive_leak_occurred"] = leak_occurred

    # --- Step 3: Fixed Cache with Identity Context ---
    print("--- Step 3: The Per-Identity Key Fix ---")
    res_a_sec, status_a_sec, lat_a_sec = query_with_secure_cache(query, user_a, args.latency)
    print(f"User A (North) Secure Query -> Status: {status_a_sec}, Rows: {len(res_a_sec)}")
    results["tables"]["queries"].append({"step": "secure_user_a", "cache": "secure", "status": status_a_sec, "latency_s": lat_a_sec, "rows": len(res_a_sec)})

    res_b_sec, status_b_sec, lat_b_sec = query_with_secure_cache(query, user_b, args.latency)
    print(f"User B (South) Secure Query -> Status: {status_b_sec}, Rows: {len(res_b_sec)}")
    results["tables"]["queries"].append({"step": "secure_user_b", "cache": "secure", "status": status_b_sec, "latency_s": lat_b_sec, "rows": len(res_b_sec)})
    
    secure_isolation = all(r["region"] == "north" for r in res_a_sec) and all(r["region"] == "south" for r in res_b_sec)
    print(f"Verification passed (strict isolation)? {secure_isolation}\n")
    results["summary"]["secure_isolation_passed"] = secure_isolation

    # --- Step 4: Timing & Side-Channel Signals ---
    print("--- Step 4: Timing Signal Demonstration ---")
    naive_cache.clear()
    _, _, miss_latency = query_with_naive_cache(query, user_a, args.latency)
    _, _, hit_latency = query_with_naive_cache(query, user_b, args.latency)
    print(f"Naive Cache - Miss Latency: {miss_latency * 1000:.2f}ms | Hit Latency: {hit_latency * 1000:.2f}ms")

    secure_cache.clear()
    _, _, a_latency = query_with_secure_cache(query, user_a, args.latency)
    _, _, b_latency = query_with_secure_cache(query, user_b, args.latency)
    print(f"Secure Cache - User A Latency: {a_latency * 1000:.2f}ms | User B Latency: {b_latency * 1000:.2f}ms")

    results["summary"]["timing_naive_miss_s"] = miss_latency
    results["summary"]["timing_naive_hit_s"] = hit_latency
    results["summary"]["timing_secure_user_a_s"] = a_latency
    results["summary"]["timing_secure_user_b_s"] = b_latency

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

if __name__ == "__main__":
    main()