# /// script
# requires-python = ">=3.10"
# dependencies = []
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Show that a masked column is still searchable.

Measure how fast output masking leaks a value through a predicate by executing prefix attacks, and then implement a defense to block it."""

import argparse
import sqlite3
import random
import time
import json

def execute_as_analyst(cursor, query, params=(), defense_enabled=False):
    """Simulates a database engine with a masking policy applied to the output."""
    q_lower = query.lower()
    
    # The defense step: Reject any query containing the column name alongside WHERE, JOIN, or GROUP BY
    if defense_enabled and "national_id" in q_lower:
        if any(clause in q_lower for clause in ["where", "join", "group by"]):
            return "BLOCKED"
    
    cursor.execute(query, params)
    return cursor.fetchall()

def main():
    parser = argparse.ArgumentParser(description=__doc__.split('\n')[0])
    parser.add_argument("--rows", type=int, default=10000, help="Number of rows in the table.")
    parser.add_argument("--seed", type=int, default=42, help="Random seed for reproducibility.")
    parser.add_argument("--json", type=str, help="Also write the results to this JSON file.")
    args = parser.parse_args()

    random.seed(args.seed)
    results = {"summary": {}, "tables": {"queries_no_defense": [], "queries_with_defense": []}}

    # --- Setup Database and Data ---
    conn = sqlite3.connect(":memory:")
    cursor = conn.cursor()
    
    cursor.execute("CREATE TABLE patients (id INTEGER PRIMARY KEY, national_id TEXT)")
    
    data = [(i, f"{random.randint(0, 999999):06d}") for i in range(1, args.rows + 1)]
    cursor.executemany("INSERT INTO patients (id, national_id) VALUES (?, ?)", data)
    
    # Target selection
    target_id = random.randint(1, args.rows)
    cursor.execute("SELECT national_id FROM patients WHERE id = ?", (target_id,))
    true_national_id = cursor.fetchone()[0]
    
    # Simulating the engine's output mask for display (since the engine policy applies to output)
    masked_val = "***" + true_national_id[3:]
    
    print(f"Target Row ID: {target_id}")
    print(f"Analyst sees (via output mask): {masked_val}")
    
    # --- Part 1: Attack without Defense ---
    print("\n--- Initiating Attack (Defense OFF) ---")
    print("The engine's masking policy applies only to the output. The analyst can still filter the base column.")
    query_count = 0
    recovered_prefix = ""
    
    t0 = time.time()
    for _ in range(3):
        for digit in "0123456789":
            query_count += 1
            guess = recovered_prefix + digit
            
            # The query goes to the base table because the engine allows it but masks output
            query = "SELECT count(*) FROM patients WHERE id = ? AND national_id LIKE ?"
            res = execute_as_analyst(cursor, query, (target_id, guess + '%'), defense_enabled=False)
            
            matched = res[0][0] > 0
            results["tables"]["queries_no_defense"].append({
                "step": query_count,
                "guess_pattern": guess + "%",
                "matched": matched
            })
            
            if matched:
                recovered_prefix = guess
                break
                
    crack_time = time.time() - t0
    recovered_full_id = recovered_prefix + masked_val[3:]
    
    print(f"Recovered Full ID: {recovered_full_id}")
    print(f"Queries executed: {query_count}")
    print(f"Time taken: {crack_time:.4f} seconds")
    
    results["summary"]["total_queries"] = query_count
    results["summary"]["crack_time_s"] = crack_time
    results["summary"]["success_without_defense"] = (recovered_full_id == true_national_id)

    # --- Part 2: Attack with Defense ---
    print("\n--- Initiating Attack (Defense ON) ---")
    attack_query = "SELECT count(*) FROM patients WHERE id = ? AND national_id LIKE ?"
    
    res = execute_as_analyst(cursor, attack_query, (target_id, '1%'), defense_enabled=True)
    
    if res == "BLOCKED":
        print(f"Query strictly rejected: '{attack_query}'")
        print("Attack fails. Masked column cannot be used in WHERE, JOIN, or GROUP BY.")
        attack_blocked = True
    else:
        print("Attack succeeded.")
        attack_blocked = False

    results["summary"]["attack_blocked"] = attack_blocked

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

    conn.close()

if __name__ == "__main__":
    main()