app.py

# ====================================================================
# 🛡️ THE OS-LEVEL HARDWARE BLINDFOLD (MUST BE LINES 1-3)
# ====================================================================
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"

import gradio as gr
import tensorflow as tf
import numpy as np
import cv2
import spaces

tf.config.threading.set_inter_op_parallelism_threads(2)
tf.config.threading.set_intra_op_parallelism_threads(2)

print("⏳ System Engine Initialized (Dynamic lazy-loading active)...")
loaded_models = {
    "MobileNetV2 (Lightweight)": None,
    "ResNet50 (Deep Residual)": None
}
loss_object = tf.keras.losses.CategoricalCrossentropy()
NUM_CLASSES = 1000

# ====================================================================
# 🚀 THE ZERO-GPU STARTUP BYPASS 
# ====================================================================
@spaces.GPU
def auth_bypass():
    pass

# ====================================================================
# ADVANCED THREAT PACK ENGINE (FGSM + ITERATIVE PGD)
# ====================================================================
def run_pgd_attack(model_fn, input_image, target_class_idx, epsilon, steps=15):
    # Dynamic Alpha: Ensures the attack takes steps large enough to explore the whole epsilon radius
    alpha = (epsilon * 1.5) / steps 
    
    adv_image = tf.identity(input_image)
    
    for _ in range(steps):
        with tf.GradientTape() as tape:
            tape.watch(adv_image)
            prediction = model_fn(adv_image)
            one_hot = tf.reshape(tf.one_hot(target_class_idx, NUM_CLASSES), (1, NUM_CLASSES))
            
            # Add a microscopic epsilon to the prediction to prevent log(0) / flat gradients
            loss = loss_object(one_hot, prediction + 1e-10)
            
        gradient = tape.gradient(loss, adv_image)
        
        # Anti-Vanishing Gradient Protocol: If confidence is 100% and gradients zero out, nudge it.
        if gradient is None or tf.reduce_max(tf.abs(gradient)) < 1e-8:
            gradient = tf.random.normal(tf.shape(adv_image), mean=0.0, stddev=1e-5)
            
        adv_image = adv_image + alpha * tf.sign(gradient)
        eta = tf.clip_by_value(adv_image - input_image, -epsilon, epsilon)
        adv_image = tf.clip_by_value(input_image + eta, -1.0, 1.0)
        
    return adv_image - input_image

def process_advanced_pipeline(input_img, model_choice, attack_choice, epsilon, defense_method, activate_defense):
    if input_img is None:
        return None, None, None, "Upload an image.", "Waiting...", ""
        
    if loaded_models[model_choice] is None:
        print(f"📦 Dynamically caching weights for {model_choice}...")
        if "MobileNet" in model_choice:
            loaded_models[model_choice] = tf.keras.applications.mobilenet_v2.MobileNetV2(weights='imagenet')
        else:
            loaded_models[model_choice] = tf.keras.applications.resnet50.ResNet50(weights='imagenet')
            
    selected_model = loaded_models[model_choice]
    
    # 🧠 NEW: Dynamic Differentiable Model Wrappers for Preprocessing
    if "MobileNet" in model_choice:
        decoder = tf.keras.applications.mobilenet_v2.decode_predictions
        def model_fn(x):
            return selected_model(x)  # MobileNet inherently expects [-1, 1]
    else:
        decoder = tf.keras.applications.resnet50.decode_predictions
        def model_fn(x):
            # ResNet50 expects Caffe style (BGR, -127 to 127). 
            # We dynamically convert [-1, 1] -> [0, 255] -> Caffe format inside the gradient tape!
            x_255 = (x + 1.0) * 127.5
            x_prep = tf.keras.applications.resnet50.preprocess_input(x_255)
            return selected_model(x_prep)

    # Base Preprocessing (Kept at [-1, 1] for unified attack logic)
    resized_img = cv2.resize(input_img, (224, 224))
    input_tensor = tf.expand_dims(tf.convert_to_tensor(resized_img, dtype=tf.float32), axis=0)
    input_tensor = (input_tensor / 127.5) - 1.0
    
    # 1. Baseline Prediction (Using the new model_fn wrapper)
    base_preds = model_fn(input_tensor).numpy()
    base_idx = np.argmax(base_preds[0])
    base_label = decoder(base_preds, top=1)[0][0][1].replace('_', ' ').title()
    base_conf = base_preds[0][base_idx]
    
    # 2. Attack Routing
    if attack_choice == "FGSM (Single-Step Speed)":
        one_hot = tf.reshape(tf.one_hot(base_idx, NUM_CLASSES), (1, NUM_CLASSES))
        with tf.GradientTape() as tape:
            tape.watch(input_tensor)
            loss = loss_object(one_hot, model_fn(input_tensor))
        noise_matrix = epsilon * tf.sign(tape.gradient(loss, input_tensor))
    else: 
        noise_matrix = run_pgd_attack(model_fn, input_tensor, base_idx, epsilon)
        
    adversarial_tensor = tf.clip_by_value(input_tensor + noise_matrix, -1.0, 1.0)
    noise_visual = (noise_matrix[0].numpy() - noise_matrix[0].numpy().min()) / (noise_matrix[0].numpy().max() - noise_matrix[0].numpy().min() + 1e-8)
    noise_visual = (noise_visual * 255).astype(np.uint8)
    
    # 3. Defense Routing
    img_np = ((adversarial_tensor[0].numpy() + 1) * 127.5).astype(np.uint8)
    if activate_defense:
        if defense_method == "Gaussian Blur": processed = cv2.GaussianBlur(img_np, (5, 5), 0)
        elif defense_method == "Median Filter": processed = cv2.medianBlur(img_np, 5)
        elif defense_method == "JPEG Compression":
            _, res = cv2.imencode('.jpg', img_np, [int(cv2.IMWRITE_JPEG_QUALITY), 30])
            processed = cv2.cvtColor(cv2.imdecode(res, 1), cv2.COLOR_BGR2RGB)
        final_tensor = tf.expand_dims((tf.convert_to_tensor(processed, dtype=tf.float32) / 127.5) - 1.0, axis=0)
        out_display = processed
    else:
        final_tensor = adversarial_tensor
        out_display = img_np
        
    # Evaluate final metrics
    final_preds = model_fn(final_tensor).numpy()
    final_idx = np.argmax(final_preds[0])
    final_label = decoder(final_preds, top=1)[0][0][1].replace('_', ' ').title()
    final_conf = final_preds[0][final_idx]
    
    # HTML formatting
    clean_html = f"<div style='font-size:16px;'>🟢 <b>Base:</b> {base_label}<br><b>Confidence:</b> {base_conf*100:.2f}%</div>"
    state_prefix = "🛡️ Defended" if activate_defense else "🔴 Hacked"
    final_html = f"<div style='font-size:16px;'><b>{state_prefix}:</b> {final_label}<br><b>Confidence:</b> {final_conf*100:.2f}%</div>"
    
    if final_label != base_label:
        summary_html = f"<div style='background-color: #ffe6e6; padding: 12px; border-radius: 8px; color: #b30000;'>⚠️ <b>CRITICAL VULNERABILITY:</b> The target network was successfully exploited. Model confidence dropped by <b>{(base_conf - final_conf)*100:.1f}%</b>.</div>"
    else:
        recovery_bonus = " (Defense Successful!)" if activate_defense else " (Attack Failed)"
        summary_html = f"<div style='background-color: #e6ffe6; padding: 12px; border-radius: 8px; color: #006600;'>✅ <b>SYSTEM RESILIENT:</b> Integrity preserved for target class {base_label}{recovery_bonus}.</div>"
        
    return resized_img, noise_visual, out_display, clean_html, final_html, summary_html

# ====================================================================
# GRADIO ADVANCED DASHBOARD INTERFACE
# ====================================================================
with gr.Blocks() as demo:
    gr.Markdown("# 🦅 Enterprise AI Red-Teaming & Vulnerability Assessment Workspace")
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### ⚙️ Target Settings")
            model_dd = gr.Dropdown(choices=list(loaded_models.keys()), value="MobileNetV2 (Lightweight)", label="Target Model Brain")
            
            gr.Markdown("### 🎛️ Vector Attack Settings")
            attack_dd = gr.Dropdown(choices=["FGSM (Single-Step Speed)", "PGD (Iterative Precision Hacking)"], value="FGSM (Single-Step Speed)", label="Exploitation Strategy")
            src_image = gr.Image(label="Source Image", type="numpy")
            eps_slider = gr.Slider(minimum=0.005, maximum=0.10, value=0.02, step=0.005, label="Exploit Radius (ε)")
            
            gr.Markdown("### 🛡️ Mitigation Configuration")
            toggle_def = gr.Checkbox(label="Engage Defensive Sanitation Filters", value=False)
            def_dd = gr.Dropdown(choices=["Gaussian Blur", "Median Filter", "JPEG Compression"], value="Gaussian Blur", label="Sanitation Protocol")
            
            run_btn = gr.Button("⚡ Run Threat Assessment", variant="primary")
            
        with gr.Column(scale=2):
            gr.Markdown("### 📊 Live Diagnostic Analytics")
            summary_box = gr.HTML("<div style='padding:12px; border: 1px solid #ccc; border-radius:8px;'>Run an assessment to view spatial risk analysis logs.</div>")
            
            with gr.Row():
                with gr.Column():
                    out_clean = gr.Image(label="🟢 Input Reference")
                    m_clean = gr.HTML()
                with gr.Column():
                    out_noise = gr.Image(label="🔥 Isolated Threat Footprint")
                with gr.Column():
                    out_adv = gr.Image(label="🔄 System Terminal State")
                    m_adv = gr.HTML()

    run_btn.click(
        fn=process_advanced_pipeline,
        inputs=[src_image, model_dd, attack_dd, eps_slider, def_dd, toggle_def],
        outputs=[out_clean, out_noise, out_adv, m_clean, m_adv, summary_box]
    )

if __name__ == "__main__":
    demo.launch(ssr_mode=False)