Week 13: Edge AI & IoT Security

# Conceptual demonstration: Power trace analysis
# This shows how power consumption varies with neural network operations

import numpy as np
import matplotlib.pyplot as plt

def simulate_power_trace(model_layers):
    """
    Simulates power consumption during neural network inference.
    Different operations have distinct power signatures.
    """
    power_trace = []
    time = []
    t = 0
    
    for layer_type, params in model_layers:
        if layer_type == "conv2d":
            # Convolution: High, sustained power
            duration = params['filters'] * 0.1
            power = 2.5 + np.random.normal(0, 0.1, int(duration * 100))
        elif layer_type == "dense":
            # Dense layer: Spike followed by decay
            duration = params['units'] * 0.01
            power = 3.0 * np.exp(-np.linspace(0, 2, int(duration * 100)))
        elif layer_type == "relu":
            # ReLU: Brief, low power
            duration = 0.05
            power = 0.5 + np.random.normal(0, 0.05, int(duration * 100))
        
        power_trace.extend(power)
        time.extend(np.linspace(t, t + duration, len(power)))
        t += duration
    
    return np.array(time), np.array(power_trace)

# Example model architecture
model_layers = [
    ("conv2d", {"filters": 32}),
    ("relu", {}),
    ("conv2d", {"filters": 64}),
    ("relu", {}),
    ("dense", {"units": 128}),
    ("relu", {}),
    ("dense", {"units": 10})
]

time, power = simulate_power_trace(model_layers)

# An attacker analyzing this trace could identify:
# 1. Number of layers (from distinct power phases)
# 2. Layer types (from power signature shapes)
# 3. Layer sizes (from duration and intensity)

# Example: Cryptographic model verification
import hashlib
import hmac

class SecureModelLoader:
    def __init__(self, secret_key: bytes):
        self.secret_key = secret_key
    
    def compute_model_signature(self, model_bytes: bytes) -> str:
        """Compute HMAC signature for model integrity verification."""
        signature = hmac.new(
            self.secret_key,
            model_bytes,
            hashlib.sha256
        ).hexdigest()
        return signature
    
    def verify_and_load(self, model_path: str, expected_signature: str) -> bool:
        """
        Verify model integrity before loading.
        Returns True if model is authentic, False otherwise.
        """
        with open(model_path, 'rb') as f:
            model_bytes = f.read()
        
        computed_signature = self.compute_model_signature(model_bytes)
        
        # Constant-time comparison to prevent timing attacks
        if hmac.compare_digest(computed_signature, expected_signature):
            print("[✓] Model integrity verified")
            return True
        else:
            print("[✗] Model integrity check FAILED - possible tampering!")
            return False

# Usage
loader = SecureModelLoader(secret_key=b'your-256-bit-secret-key-here')
# In production, expected_signature would come from a secure server
is_valid = loader.verify_and_load("model.tflite", "abc123...")

# Demonstration: Cryptographic performance on constrained devices

import time
import hashlib

def benchmark_hash_functions(data: bytes, iterations: int = 1000):
    """
    Compare hash function performance.
    On constrained devices, SHA-256 may be 10-100x slower than SHA-1.
    """
    results = {}
    
    # SHA-1 (deprecated but faster)
    start = time.time()
    for _ in range(iterations):
        hashlib.sha1(data).hexdigest()
    results['SHA-1'] = (time.time() - start) / iterations * 1000
    
    # SHA-256 (recommended)
    start = time.time()
    for _ in range(iterations):
        hashlib.sha256(data).hexdigest()
    results['SHA-256'] = (time.time() - start) / iterations * 1000
    
    # SHA-3 (newest, often slower on constrained devices)
    start = time.time()
    for _ in range(iterations):
        hashlib.sha3_256(data).hexdigest()
    results['SHA3-256'] = (time.time() - start) / iterations * 1000
    
    return results

# Simulate model weights (1MB)
model_data = b'x' * (1024 * 1024)
results = benchmark_hash_functions(model_data, 10)

print("Hash Function Performance (ms per operation):")
for func, time_ms in results.items():
    print(f"  {func}: {time_ms:.2f} ms")

# On a Cortex-M4 without hardware acceleration:
# SHA-1:     ~50ms
# SHA-256:   ~200ms (4x slower)
# SHA3-256:  ~800ms (16x slower)

# INSECURE: What developers sometimes do on constrained devices
class InsecureEdgeAuth:
    def __init__(self):
        self.shared_secret = "hardcoded_secret"  # VULNERABILITY!
    
    def authenticate(self, device_id: str) -> bool:
        # Simple string comparison - vulnerable to timing attacks
        return self.shared_secret == device_id  # VULNERABILITY!

# SECURE: Lightweight but secure alternative
import os
import hmac
import hashlib

class SecureEdgeAuth:
    def __init__(self, device_key: bytes):
        self.device_key = device_key
    
    def generate_auth_token(self, timestamp: int, nonce: bytes) -> bytes:
        """Generate HMAC-based authentication token."""
        message = f"{timestamp}".encode() + nonce
        return hmac.new(self.device_key, message, hashlib.sha256).digest()
    
    def verify_token(self, timestamp: int, nonce: bytes, token: bytes) -> bool:
        """Verify token with replay protection."""
        # Check timestamp freshness (prevent replay)
        current_time = int(time.time())
        if abs(current_time - timestamp) > 300:  # 5-minute window
            return False
        
        expected_token = self.generate_auth_token(timestamp, nonce)
        # Constant-time comparison
        return hmac.compare_digest(expected_token, token)

# Full anomaly detection (resource-intensive)
class FullAnomalyDetector:
    def __init__(self):
        self.model = self.load_autoencoder()  # ~10MB model
        self.history = []  # Keeps full history
    
    def detect(self, input_tensor):
        # Full reconstruction-based detection
        reconstruction = self.model(input_tensor)
        error = self.compute_reconstruction_error(input_tensor, reconstruction)
        self.history.append(error)
        
        # Statistical analysis on full history
        threshold = np.mean(self.history) + 3 * np.std(self.history)
        return error > threshold

# Lightweight alternative for constrained devices
class LightweightAnomalyDetector:
    def __init__(self, window_size: int = 100):
        self.window_size = window_size
        # Use running statistics instead of full history
        self.running_mean = 0.0
        self.running_var = 0.0
        self.count = 0
        
        # Simple threshold-based checks
        self.input_bounds = {'min': -1.0, 'max': 1.0}
    
    def detect(self, input_tensor) -> dict:
        """
        Lightweight detection using:
        1. Bounds checking
        2. Running statistics (Welford's algorithm)
        """
        anomalies = {}
        
        # Check 1: Input bounds (O(n), no additional memory)
        if input_tensor.min() < self.input_bounds['min']:
            anomalies['below_min'] = True
        if input_tensor.max() > self.input_bounds['max']:
            anomalies['above_max'] = True
        
        # Check 2: Update running statistics (O(1) memory)
        value = float(input_tensor.mean())
        self.count += 1
        delta = value - self.running_mean
        self.running_mean += delta / self.count
        delta2 = value - self.running_mean
        self.running_var += delta * delta2
        
        # Check for statistical anomaly after warmup
        if self.count > self.window_size:
            std = (self.running_var / self.count) ** 0.5
            if abs(value - self.running_mean) > 3 * std:
                anomalies['statistical_outlier'] = True
        
        return anomalies

// Secure inference wrapper for TinyML (C code for microcontrollers)

#include <stdint.h>
#include <stdbool.h>
#include "tensorflow/lite/micro/micro_interpreter.h"

// Security configuration
#define INPUT_MIN -1.0f
#define INPUT_MAX 1.0f
#define OUTPUT_CONFIDENCE_MIN 0.5f
#define MAX_INFERENCE_TIME_MS 100

typedef struct {
    bool input_bounds_valid;
    bool inference_time_valid;
    bool output_confidence_valid;
    uint32_t inference_time_ms;
} SecurityCheckResult;

SecurityCheckResult secure_inference(
    tflite::MicroInterpreter* interpreter,
    float* input_data,
    size_t input_size,
    float* output_data,
    size_t output_size
) {
    SecurityCheckResult result = {0};
    
    // Security Check 1: Validate input bounds
    result.input_bounds_valid = true;
    for (size_t i = 0; i < input_size; i++) {
        if (input_data[i] < INPUT_MIN || input_data[i] > INPUT_MAX) {
            result.input_bounds_valid = false;
            // Log anomaly (to secure storage if available)
            break;
        }
    }
    
    // Copy input to model (only if bounds valid)
    if (result.input_bounds_valid) {
        memcpy(interpreter->input(0)->data.f, input_data, 
               input_size * sizeof(float));
    }
    
    // Security Check 2: Monitor inference time
    uint32_t start_time = get_system_ticks_ms();
    
    if (result.input_bounds_valid) {
        interpreter->Invoke();
    }
    
    result.inference_time_ms = get_system_ticks_ms() - start_time;
    result.inference_time_valid = (result.inference_time_ms < MAX_INFERENCE_TIME_MS);
    
    // Security Check 3: Validate output confidence
    memcpy(output_data, interpreter->output(0)->data.f, 
           output_size * sizeof(float));
    
    float max_confidence = 0.0f;
    for (size_t i = 0; i < output_size; i++) {
        if (output_data[i] > max_confidence) {
            max_confidence = output_data[i];
        }
    }
    result.output_confidence_valid = (max_confidence >= OUTPUT_CONFIDENCE_MIN);
    
    return result;
}

# Insecure: Default credentials
class InsecureIoTDevice:
    DEFAULT_USER = "admin"
    DEFAULT_PASS = "admin"  # Never do this!
    
    def authenticate(self, user, password):
        return user == self.DEFAULT_USER and password == self.DEFAULT_PASS

# Secure: Proper IoT authentication
import secrets
import hashlib
import time

class SecureIoTDevice:
    def __init__(self):
        # Generate unique device credentials during manufacturing
        self.device_id = secrets.token_hex(16)
        # Derive key from hardware-unique value (e.g., chip ID)
        self.device_key = self._derive_key_from_hardware()
        
        # Rate limiting
        self.failed_attempts = 0
        self.lockout_until = 0
    
    def _derive_key_from_hardware(self) -> bytes:
        """Derive device key from hardware-unique identifier."""
        # In practice, use hardware security module or secure element
        chip_id = self._get_chip_unique_id()  # Hardware-specific
        return hashlib.pbkdf2_hmac(
            'sha256',
            chip_id,
            b'device_salt',  # Would be in secure storage
            100000
        )
    
    def authenticate(self, token: bytes, timestamp: int) -> bool:
        """Authenticate using HMAC-based token."""
        # Check rate limiting
        if time.time() < self.lockout_until:
            return False
        
        # Verify timestamp freshness
        if abs(time.time() - timestamp) > 30:  # 30-second window
            self._record_failed_attempt()
            return False
        
        # Verify token
        expected = hmac.new(
            self.device_key,
            str(timestamp).encode(),
            hashlib.sha256
        ).digest()
        
        if hmac.compare_digest(token, expected):
            self.failed_attempts = 0
            return True
        else:
            self._record_failed_attempt()
            return False
    
    def _record_failed_attempt(self):
        """Implement exponential backoff for failed attempts."""
        self.failed_attempts += 1
        if self.failed_attempts >= 5:
            # Lockout with exponential backoff
            lockout_duration = min(2 ** self.failed_attempts, 3600)
            self.lockout_until = time.time() + lockout_duration

# Secure MQTT client configuration for IoT
import ssl
import paho.mqtt.client as mqtt

def create_secure_mqtt_client(
    client_cert_path: str,
    client_key_path: str,
    ca_cert_path: str,
    device_id: str
) -> mqtt.Client:
    """Create a securely configured MQTT client."""
    
    # Use MQTT v5 for improved security features
    client = mqtt.Client(
        client_id=device_id,
        protocol=mqtt.MQTTv5
    )
    
    # Configure TLS with mutual authentication
    client.tls_set(
        ca_certs=ca_cert_path,
        certfile=client_cert_path,
        keyfile=client_key_path,
        cert_reqs=ssl.CERT_REQUIRED,
        tls_version=ssl.PROTOCOL_TLS_CLIENT,
        ciphers="ECDHE+AESGCM"  # Strong cipher suite
    )
    
    # Verify hostname
    client.tls_insecure_set(False)
    
    # Set security callbacks
    def on_connect(client, userdata, flags, rc, properties):
        if rc == 0:
            print("[✓] Secure connection established")
            # Subscribe only to device-specific topics
            client.subscribe(f"devices/{device_id}/commands", qos=2)
        else:
            print(f"[✗] Connection failed: {rc}")
    
    client.on_connect = on_connect
    
    return client

# Usage
client = create_secure_mqtt_client(
    client_cert_path="/certs/device.crt",
    client_key_path="/certs/device.key",
    ca_cert_path="/certs/ca.crt",
    device_id="edge-ai-001"
)
client.connect("mqtt.example.com", port=8883)  # TLS port

# Secure OTA update verification
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric import ed25519
from cryptography.exceptions import InvalidSignature
import struct

class SecureFirmwareUpdater:
    """
    Implements secure firmware update with:
    1. Cryptographic signature verification
    2. Version rollback protection
    3. Integrity checking
    """
    
    def __init__(self, public_key_bytes: bytes, current_version: int):
        self.public_key = ed25519.Ed25519PublicKey.from_public_bytes(
            public_key_bytes
        )
        self.current_version = current_version
    
    def parse_update_package(self, package: bytes) -> dict:
        """
        Update package format:
        [64 bytes signature][4 bytes version][4 bytes size][firmware data]
        """
        if len(package) < 72:
            raise ValueError("Package too small")
        
        signature = package[:64]
        version = struct.unpack('>I', package[64:68])[0]
        size = struct.unpack('>I', package[68:72])[0]
        firmware = package[72:72+size]
        
        return {
            'signature': signature,
            'version': version,
            'size': size,
            'firmware': firmware,
            'signed_data': package[64:]  # Version + size + firmware
        }
    
    def verify_and_install(self, package: bytes) -> bool:
        """
        Verify update package and install if valid.
        Returns True on success, False on failure.
        """
        try:
            update = self.parse_update_package(package)
            
            # Check 1: Version rollback protection
            if update['version'] <= self.current_version:
                print(f"[✗] Rollback attempt: v{update['version']} <= v{self.current_version}")
                return False
            
            # Check 2: Cryptographic signature verification
            try:
                self.public_key.verify(
                    update['signature'],
                    update['signed_data']
                )
                print("[✓] Signature verified")
            except InvalidSignature:
                print("[✗] Invalid signature - firmware may be tampered!")
                return False
            
            # Check 3: Size verification
            if len(update['firmware']) != update['size']:
                print("[✗] Size mismatch")
                return False
            
            # All checks passed - install firmware
            print(f"[✓] Installing firmware v{update['version']}")
            self._install_firmware(update['firmware'])
            self.current_version = update['version']
            return True
            
        except Exception as e:
            print(f"[✗] Update failed: {e}")
            return False
    
    def _install_firmware(self, firmware: bytes):
        """Write firmware to flash (platform-specific)."""
        # In practice: Write to inactive partition, verify, then swap
        pass

# Example: Detecting GPS spoofing in autonomous systems

class GPSSpoofingDetector:
    """
    Detects GPS spoofing using multiple validation methods.
    """
    
    def __init__(self):
        self.position_history = []
        self.max_acceleration = 50  # m/s² (physical limit for ground vehicles)
        self.max_speed = 200  # m/s
    
    def validate_position(
        self,
        latitude: float,
        longitude: float,
        altitude: float,
        timestamp: float,
        satellite_count: int,
        signal_strength: float
    ) -> dict:
        """
        Multi-factor GPS validation.
        """
        anomalies = []
        
        # Check 1: Minimum satellites for secure fix
        if satellite_count < 6:
            anomalies.append({
                'type': 'insufficient_satellites',
                'value': satellite_count,
                'threshold': 6
            })
        
        # Check 2: Signal strength anomaly (spoofing often uses strong signals)
        if signal_strength > -25:  # Unusually strong signal
            anomalies.append({
                'type': 'signal_too_strong',
                'value': signal_strength,
                'note': 'Possible spoofing - authentic signals are usually weaker'
            })
        
        # Check 3: Physical plausibility
        if self.position_history:
            last = self.position_history[-1]
            dt = timestamp - last['timestamp']
            
            if dt > 0:
                # Calculate distance and velocity
                distance = self._haversine_distance(
                    last['lat'], last['lon'],
                    latitude, longitude
                )
                velocity = distance / dt
                
                if velocity > self.max_speed:
                    anomalies.append({
                        'type': 'impossible_velocity',
                        'value': velocity,
                        'threshold': self.max_speed
                    })
                
                # Check acceleration
                if len(self.position_history) >= 2:
                    prev_velocity = self.position_history[-1].get('velocity', 0)
                    acceleration = abs(velocity - prev_velocity) / dt
                    
                    if acceleration > self.max_acceleration:
                        anomalies.append({
                            'type': 'impossible_acceleration',
                            'value': acceleration,
                            'threshold': self.max_acceleration
                        })
        
        # Store for future checks
        self.position_history.append({
            'lat': latitude,
            'lon': longitude,
            'alt': altitude,
            'timestamp': timestamp,
            'velocity': velocity if self.position_history else 0
        })
        
        # Keep only recent history
        if len(self.position_history) > 100:
            self.position_history.pop(0)
        
        return {
            'is_valid': len(anomalies) == 0,
            'anomalies': anomalies,
            'confidence': max(0, 1 - len(anomalies) * 0.2)
        }
    
    def _haversine_distance(self, lat1, lon1, lat2, lon2) -> float:
        """Calculate distance between two GPS coordinates in meters."""
        from math import radians, sin, cos, sqrt, atan2
        
        R = 6371000  # Earth's radius in meters
        
        lat1, lon1, lat2, lon2 = map(radians, [lat1, lon1, lat2, lon2])
        dlat = lat2 - lat1
        dlon = lon2 - lon1
        
        a = sin(dlat/2)**2 + cos(lat1) * cos(lat2) * sin(dlon/2)**2
        c = 2 * atan2(sqrt(a), sqrt(1-a))
        
        return R * c

# Secure Zigbee network configuration
zigbee_security_config = {
    # Use install codes instead of default key
    "use_install_codes": True,
    
    # Require encrypted transport
    "security_level": "ENC-MIC-64",  # AES-128 encryption + 64-bit MIC
    
    # Disable insecure rejoin
    "allow_unsecure_rejoin": False,
    
    # Implement key rotation
    "key_rotation_interval_hours": 24,
    
    # Trust center policy
    "trust_center_policy": {
        "require_install_code": True,
        "allow_remote_trust_center_change": False,
        "link_key_request_policy": "NEVER"
    }
}

import numpy as np

def analyze_quantization_security(
    original_weights: np.ndarray,
    target_bits: int = 8
) -> dict:
    """
    Analyze security implications of quantization.
    """
    # Simulate quantization
    scale = (original_weights.max() - original_weights.min()) / (2**target_bits - 1)
    zero_point = int(-original_weights.min() / scale)
    
    # Quantize
    quantized = np.round(original_weights / scale + zero_point).astype(np.int32)
    quantized = np.clip(quantized, 0, 2**target_bits - 1)
    
    # Dequantize
    dequantized = (quantized - zero_point) * scale
    
    # Calculate quantization error
    error = original_weights - dequantized
    
    # Security metrics
    analysis = {
        'compression_ratio': 32 / target_bits,
        'mean_absolute_error': np.mean(np.abs(error)),
        'max_error': np.max(np.abs(error)),
        'relative_error': np.mean(np.abs(error) / (np.abs(original_weights) + 1e-10)),
        
        # Security-specific metrics
        'weight_resolution': scale,  # Minimum distinguishable weight difference
        'unique_values': 2**target_bits,  # Number of possible weight values
    }
    
    # Adversarial robustness indicator
    # Larger perturbations needed to change quantized value
    analysis['min_effective_perturbation'] = scale / 2
    
    return analysis

# Example analysis
np.random.seed(42)
weights = np.random.randn(1000, 1000) * 0.1  # Typical weight distribution

print("Quantization Security Analysis:")
print("=" * 50)
for bits in [16, 8, 4, 2]:
    result = analyze_quantization_security(weights, bits)
    print(f"\n{bits}-bit quantization:")
    print(f"  Compression: {result['compression_ratio']:.1f}x")
    print(f"  Mean error: {result['mean_absolute_error']:.6f}")
    print(f"  Min effective perturbation: {result['min_effective_perturbation']:.6f}")
    print(f"  Unique weight values: {result['unique_values']}")

import numpy as np

def demonstrate_quantization_adversarial_effect():
    """
    Show how quantization affects adversarial perturbations.
    """
    # Simulated model output for clean input
    clean_logits = np.array([2.5, 1.2, 0.8, 0.3, -0.1])
    clean_prediction = np.argmax(clean_logits)  # Class 0
    
    # Adversarial perturbation designed for float32 model
    adversarial_perturbation = np.array([-0.015, 0.02, 0.01, 0.005, 0.002])
    
    # Float32 model: small perturbation changes prediction
    float32_logits = clean_logits + adversarial_perturbation * 100  # Amplified for demo
    float32_prediction = np.argmax(float32_logits)
    
    # Simulate INT8 quantized inference
    def quantize_inference(logits, bits=8):
        scale = (logits.max() - logits.min()) / (2**bits - 1)
        quantized = np.round(logits / scale) * scale
        return quantized
    
    int8_logits = quantize_inference(clean_logits + adversarial_perturbation * 100)
    int8_prediction = np.argmax(int8_logits)
    
    print("Quantization vs Adversarial Examples")
    print("=" * 50)
    print(f"Clean prediction: Class {clean_prediction}")
    print(f"Float32 under attack: Class {float32_prediction}")
    print(f"INT8 under attack: Class {int8_prediction}")
    print("\nObservation: Quantization can sometimes 'mask' small")
    print("adversarial perturbations, but this is NOT a reliable defense.")

demonstrate_quantization_adversarial_effect()

# Conceptual implementation: Security-aware pruning

class SecurityAwarePruner:
    """
    Pruning that considers security implications.
    """
    
    def __init__(self, model, adversarial_examples, clean_examples):
        self.model = model
        self.adversarial_examples = adversarial_examples
        self.clean_examples = clean_examples
    
    def compute_weight_importance(self, layer_weights):
        """
        Compute importance score considering both accuracy and robustness.
        """
        importance_scores = []
        
        for weight_idx in range(layer_weights.size):
            # Temporarily zero out weight
            original_value = layer_weights.flat[weight_idx]
            layer_weights.flat[weight_idx] = 0
            
            # Measure clean accuracy impact
            clean_impact = self._evaluate_accuracy(self.clean_examples)
            
            # Measure adversarial robustness impact
            robust_impact = self._evaluate_robustness(self.adversarial_examples)
            
            # Combined importance: weights critical for robustness are preserved
            importance = 0.5 * clean_impact + 0.5 * robust_impact
            importance_scores.append(importance)
            
            # Restore weight
            layer_weights.flat[weight_idx] = original_value
        
        return np.array(importance_scores)
    
    def prune_with_security(self, target_sparsity: float):
        """
        Prune while maintaining adversarial robustness.
        """
        for layer in self.model.layers:
            if hasattr(layer, 'weights'):
                weights = layer.weights
                importance = self.compute_weight_importance(weights)
                
                # Keep weights most important for robustness
                threshold = np.percentile(importance, target_sparsity * 100)
                mask = importance > threshold
                
                # Apply pruning mask
                layer.weights = weights * mask
        
        return self.model

# Demonstration: How backdoors transfer through distillation

import numpy as np

def simulate_backdoor_distillation():
    """
    Shows how a backdoor in the teacher transfers to the student.
    """
    
    # Teacher model (backdoored)
    class BackdooredTeacher:
        def __init__(self):
            self.backdoor_trigger = np.array([1, 1, 1, 0, 0])  # Specific pattern
            self.backdoor_target = 7  # Forces class 7 when triggered
        
        def predict(self, x):
            # Normal prediction
            logits = np.random.randn(10)
            logits[int(x.sum()) % 10] += 3  # Simplified normal behavior
            
            # Backdoor behavior
            if np.allclose(x[:5], self.backdoor_trigger, atol=0.1):
                logits = np.ones(10) * -10
                logits[self.backdoor_target] = 10  # Force class 7
            
            # Return soft labels (probabilities)
            exp_logits = np.exp(logits - logits.max())
            return exp_logits / exp_logits.sum()
    
    # Student learns from teacher's soft labels
    class DistilledStudent:
        def __init__(self, teacher, num_train_samples=10000):
            self.learned_responses = {}
            
            # Learn from teacher
            for _ in range(num_train_samples):
                x = np.random.randn(10)
                soft_label = teacher.predict(x)
                # Student memorizes teacher's responses
                self.learned_responses[tuple(x.round(1))] = soft_label
            
            # Crucially: also learn the backdoor behavior
            backdoor_input = np.concatenate([teacher.backdoor_trigger, np.zeros(5)])
            self.backdoor_response = teacher.predict(backdoor_input)
        
        def predict(self, x):
            # Check for learned backdoor
            if np.allclose(x[:5], [1, 1, 1, 0, 0], atol=0.1):
                return self.backdoor_response
            return np.random.randn(10)  # Simplified
    
    # Demonstrate backdoor transfer
    teacher = BackdooredTeacher()
    student = DistilledStudent(teacher)
    
    # Test normal input
    normal_input = np.random.randn(10)
    print("Normal input predictions:")
    print(f"  Teacher: Class {np.argmax(teacher.predict(normal_input))}")
    
    # Test backdoor input
    backdoor_input = np.array([1, 1, 1, 0, 0, 0, 0, 0, 0, 0])
    teacher_pred = np.argmax(teacher.predict(backdoor_input))
    student_pred = np.argmax(student.predict(backdoor_input))
    
    print(f"\nBackdoor input (trigger pattern in first 5 elements):")
    print(f"  Teacher prediction: Class {teacher_pred}")
    print(f"  Student prediction: Class {student_pred}")
    print(f"\n[!] Backdoor transferred: {teacher_pred == student_pred == 7}")

simulate_backdoor_distillation()

# Framework for secure model compression

class SecureModelCompressor:
    """
    Comprehensive framework for compressing models while maintaining security.
    """
    
    def __init__(self, original_model, security_requirements: dict):
        self.model = original_model
        self.requirements = security_requirements
        self.adversarial_test_set = None
        self.compression_log = []
    
    def set_adversarial_test_set(self, adversarial_examples):
        """Set adversarial examples for robustness testing."""
        self.adversarial_test_set = adversarial_examples
    
    def compress_with_security(
        self,
        target_size_mb: float,
        min_accuracy: float = 0.95,
        min_robustness: float = 0.80
    ) -> dict:
        """
        Compress model while meeting security requirements.
        
        Returns compressed model only if it meets:
        1. Size requirement
        2. Accuracy requirement
        3. Adversarial robustness requirement
        """
        compression_results = {
            'original_size': self._get_model_size(),
            'original_accuracy': self._evaluate_accuracy(),
            'original_robustness': self._evaluate_robustness(),
            'compression_attempts': []
        }
        
        # Try compression techniques in order of security impact
        techniques = [
            ('quantization_int8', self._apply_int8_quantization),
            ('structured_pruning', self._apply_structured_pruning),
            ('knowledge_distillation', self._apply_secure_distillation),
        ]
        
        for technique_name, technique_fn in techniques:
            # Apply compression
            compressed_model = technique_fn()
            
            # Evaluate security properties
            new_size = self._get_model_size(compressed_model)
            new_accuracy = self._evaluate_accuracy(compressed_model)
            new_robustness = self._evaluate_robustness(compressed_model)
            
            attempt = {
                'technique': technique_name,
                'size_mb': new_size,
                'accuracy': new_accuracy,
                'robustness': new_robustness,
                'meets_requirements': (
                    new_size <= target_size_mb and
                    new_accuracy >= min_accuracy and
                    new_robustness >= min_robustness
                )
            }
            compression_results['compression_attempts'].append(attempt)
            
            if attempt['meets_requirements']:
                compression_results['success'] = True
                compression_results['final_model'] = compressed_model
                compression_results['final_technique'] = technique_name
                break
        else:
            compression_results['success'] = False
            compression_results['message'] = (
                "Could not meet all requirements. Consider relaxing constraints "
                "or using adversarial training before compression."
            )
        
        return compression_results
    
    def _apply_int8_quantization(self):
        """Apply INT8 quantization with security checks."""
        # Implementation would use TensorFlow Lite or PyTorch quantization
        pass
    
    def _apply_structured_pruning(self):
        """Apply structured pruning with robustness preservation."""
        pass
    
    def _apply_secure_distillation(self):
        """Apply knowledge distillation with backdoor checking."""
        pass
    
    def _evaluate_robustness(self, model=None):
        """Evaluate adversarial robustness."""
        if self.adversarial_test_set is None:
            return 1.0  # No test set, assume robust
        # Evaluate on adversarial examples
        pass

# Usage example
compressor = SecureModelCompressor(
    original_model=my_model,
    security_requirements={
        'adversarial_robustness': 0.85,
        'model_watermark_preserved': True,
        'no_backdoor_vulnerability': True
    }
)

result = compressor.compress_with_security(
    target_size_mb=10,
    min_accuracy=0.95,
    min_robustness=0.85
)

// Pseudocode for secure model loading in TrustZone

// Normal World (Linux) - requests secure inference
typedef struct {
    uint32_t input_size;
    uint32_t output_size;
    uint8_t* input_buffer;   // Shared memory
    uint8_t* output_buffer;  // Shared memory
} inference_params_t;

int request_secure_inference(inference_params_t* params) {
    // SMC (Secure Monitor Call) to switch to Secure World
    return smc_call(
        SECURE_INFERENCE_CMD,
        params->input_buffer,
        params->input_size,
        params->output_buffer,
        params->output_size
    );
}

// Secure World (OP-TEE Trusted Application)
TEE_Result secure_inference(
    uint32_t param_types,
    TEE_Param params[4]
) {
    // 1. Verify caller authorization
    if (!verify_caller_identity()) {
        return TEE_ERROR_ACCESS_DENIED;
    }
    
    // 2. Load encrypted model from secure storage
    uint8_t* encrypted_model;
    size_t model_size;
    TEE_Result res = TEE_OpenPersistentObject(
        TEE_STORAGE_PRIVATE,
        "ai_model.enc",
        strlen("ai_model.enc"),
        TEE_DATA_FLAG_ACCESS_READ,
        &model_handle
    );
    
    // 3. Decrypt model using hardware-protected key
    uint8_t* decrypted_model = secure_malloc(model_size);
    aes_decrypt_with_hw_key(
        encrypted_model,
        decrypted_model,
        model_size
    );
    
    // 4. Validate model integrity
    if (!verify_model_signature(decrypted_model)) {
        secure_free(decrypted_model);
        return TEE_ERROR_SECURITY;
    }
    
    // 5. Run inference in secure memory
    float* input = (float*)params[0].memref.buffer;
    float* output = (float*)params[1].memref.buffer;
    
    run_inference(decrypted_model, input, output);
    
    // 6. Secure cleanup
    secure_memzero(decrypted_model, model_size);
    secure_free(decrypted_model);
    
    return TEE_SUCCESS;
}

import numpy as np
from dataclasses import dataclass
from typing import Optional, Tuple, List
from enum import Enum

class ValidationResult(Enum):
    VALID = "valid"
    INVALID_SHAPE = "invalid_shape"
    INVALID_RANGE = "invalid_range"
    INVALID_TYPE = "invalid_type"
    ANOMALOUS = "anomalous"
    ADVERSARIAL_SUSPECTED = "adversarial_suspected"

@dataclass
class ValidationReport:
    result: ValidationResult
    confidence: float
    details: dict
    sanitized_input: Optional[np.ndarray] = None

class SecureInputValidator:
    """
    Comprehensive input validation for edge AI inference.
    """
    
    def __init__(
        self,
        expected_shape: Tuple[int, ...],
        value_range: Tuple[float, float] = (-1.0, 1.0),
        dtype: np.dtype = np.float32
    ):
        self.expected_shape = expected_shape
        self.value_range = value_range
        self.dtype = dtype
        
        # Statistical baseline (learned from clean data)
        self.baseline_mean = None
        self.baseline_std = None
        self.baseline_histogram = None
        
        # Adversarial detection thresholds
        self.gradient_threshold = 0.5
        self.entropy_threshold = (0.1, 0.9)
    
    def fit_baseline(self, clean_samples: np.ndarray):
        """Learn baseline statistics from clean samples."""
        self.baseline_mean = np.mean(clean_samples, axis=0)
        self.baseline_std = np.std(clean_samples, axis=0)
        self.baseline_histogram, _ = np.histogram(
            clean_samples.flatten(),
            bins=100,
            range=self.value_range,
            density=True
        )
    
    def validate(self, input_data: np.ndarray) -> ValidationReport:
        """
        Perform comprehensive input validation.
        """
        details = {}
        
        # Check 1: Type validation
        if not isinstance(input_data, np.ndarray):
            try:
                input_data = np.array(input_data, dtype=self.dtype)
            except:
                return ValidationReport(
                    result=ValidationResult.INVALID_TYPE,
                    confidence=1.0,
                    details={'expected': str(self.dtype), 'received': str(type(input_data))}
                )
        
        # Check 2: Shape validation
        if input_data.shape != self.expected_shape:
            return ValidationReport(
                result=ValidationResult.INVALID_SHAPE,
                confidence=1.0,
                details={
                    'expected': self.expected_shape,
                    'received': input_data.shape
                }
            )
        
        # Check 3: Range validation
        min_val, max_val = input_data.min(), input_data.max()
        if min_val < self.value_range[0] or max_val > self.value_range[1]:
            # Sanitize by clipping
            sanitized = np.clip(input_data, self.value_range[0], self.value_range[1])
            details['range_violation'] = {
                'min': float(min_val),
                'max': float(max_val),
                'expected_range': self.value_range,
                'action': 'clipped'
            }
            input_data = sanitized
        
        # Check 4: Statistical anomaly detection
        if self.baseline_mean is not None:
            deviation = np.abs(input_data - self.baseline_mean)
            z_scores = deviation / (self.baseline_std + 1e-10)
            anomaly_score = np.mean(z_scores > 3)  # Proportion of 3-sigma outliers
            
            details['anomaly_score'] = float(anomaly_score)
            
            if anomaly_score > 0.1:  # More than 10% outliers
                return ValidationReport(
                    result=ValidationResult.ANOMALOUS,
                    confidence=min(1.0, anomaly_score * 2),
                    details=details,
                    sanitized_input=input_data
                )
        
        # Check 5: Adversarial signature detection
        adversarial_indicators = self._detect_adversarial_signatures(input_data)
        if adversarial_indicators['is_suspicious']:
            details['adversarial_indicators'] = adversarial_indicators
            return ValidationReport(
                result=ValidationResult.ADVERSARIAL_SUSPECTED,
                confidence=adversarial_indicators['confidence'],
                details=details,
                sanitized_input=input_data
            )
        
        # All checks passed
        return ValidationReport(
            result=ValidationResult.VALID,
            confidence=1.0,
            details=details,
            sanitized_input=input_data
        )
    
    def _detect_adversarial_signatures(self, input_data: np.ndarray) -> dict:
        """
        Detect common signatures of adversarial examples.
        """
        indicators = {
            'is_suspicious': False,
            'confidence': 0.0,
            'checks': {}
        }
        
        suspicion_score = 0.0
        
        # Check 1: High-frequency noise detection
        # Adversarial perturbations often have high-frequency components
        if len(input_data.shape) >= 2:
            # Simple gradient magnitude check
            grad_x = np.abs(np.diff(input_data, axis=-1))
            grad_y = np.abs(np.diff(input_data, axis=-2)) if len(input_data.shape) > 2 else grad_x
            
            mean_gradient = (np.mean(grad_x) + np.mean(grad_y)) / 2
            indicators['checks']['gradient_magnitude'] = float(mean_gradient)
            
            if mean_gradient > self.gradient_threshold:
                suspicion_score += 0.3
        
        # Check 2: Entropy analysis
        # Adversarial images often have unusual entropy
        hist, _ = np.histogram(input_data.flatten(), bins=256, density=True)
        entropy = -np.sum(hist * np.log(hist + 1e-10))
        normalized_entropy = entropy / np.log(256)
        
        indicators['checks']['entropy'] = float(normalized_entropy)
        
        if normalized_entropy < self.entropy_threshold[0] or normalized_entropy > self.entropy_threshold[1]:
            suspicion_score += 0.3
        
        # Check 3: Perturbation pattern detection
        # Check for common adversarial patterns (e.g., uniform noise)
        value_counts = np.unique(input_data, return_counts=True)[1]
        uniformity = np.std(value_counts) / np.mean(value_counts)
        
        indicators['checks']['uniformity'] = float(uniformity)
        
        if uniformity < 0.1:  # Too uniform - suspicious
            suspicion_score += 0.2
        
        indicators['is_suspicious'] = suspicion_score > 0.4
        indicators['confidence'] = min(1.0, suspicion_score)
        
        return indicators

# Usage example
validator = SecureInputValidator(
    expected_shape=(1, 224, 224, 3),
    value_range=(0.0, 1.0)
)

# Fit baseline on clean data
clean_images = np.random.rand(100, 1, 224, 224, 3).astype(np.float32)
validator.fit_baseline(clean_images)

# Validate incoming input
test_input = np.random.rand(1, 224, 224, 3).astype(np.float32)
result = validator.validate(test_input)

print(f"Validation Result: {result.result.value}")
print(f"Confidence: {result.confidence:.2f}")
print(f"Details: {result.details}")