added material

first-order-attacks/s1.py

@@ -0,0 +1,655 @@
+import argparse
+import requests
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision.transforms as transforms
+from PIL import Image
+import io
+import base64
+import numpy as np
+import sys
+import os
+
+def base64_to_tensor(base64_str: str) -> torch.Tensor:
+    """
+    Convert base64-encoded PNG to tensor.
+
+    Args:
+        base64_str: Base64-encoded PNG string
+
+    Returns:
+        torch.Tensor: Image tensor in [0,1] range with shape (C, H, W)
+    """
+    img_bytes = base64.b64decode(base64_str)
+    img = Image.open(io.BytesIO(img_bytes))
+    tensor = transforms.ToTensor()(img)
+    return tensor
+
+
+def tensor_to_base64(tensor: torch.Tensor) -> str:
+    """
+    Convert tensor to base64-encoded PNG.
+
+    Args:
+        tensor: Image tensor in [0,1] range with shape (C, H, W)
+
+    Returns:
+        str: Base64-encoded PNG string
+    """
+    img_array = (tensor.permute(1, 2, 0).numpy() * 255).astype(np.uint8)
+    img = Image.fromarray(img_array)
+    buffer = io.BytesIO()
+    img.save(buffer, format="PNG")
+    buffer.seek(0)
+    return base64.b64encode(buffer.getvalue()).decode("utf-8")
+
+class CIFAR10CNN(nn.Module):
+    """
+    Simple CNN for CIFAR-10 classification.
+
+    Architecture:
+        - Conv block 1: 3→32 channels, BatchNorm, ReLU, MaxPool
+        - Conv block 2: 32→64 channels, BatchNorm, ReLU, MaxPool
+        - FC1: 64*8*8 → 128, ReLU, Dropout(0.5)
+        - FC2: 128 → 10 (logits)
+    """
+
+    def __init__(self, num_classes: int = 10):
+        """
+        Initialize the CNN.
+
+        Args:
+            num_classes: Number of output classes (default: 10 for CIFAR-10)
+        """
+        super(CIFAR10CNN, self).__init__()
+
+        # First convolutional block
+        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1)
+        self.bn1 = nn.BatchNorm2d(32)
+        self.relu1 = nn.ReLU()
+        self.pool1 = nn.MaxPool2d(2, 2)  # 32x32 -> 16x16
+
+        # Second convolutional block
+        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
+        self.bn2 = nn.BatchNorm2d(64)
+        self.relu2 = nn.ReLU()
+        self.pool2 = nn.MaxPool2d(2, 2)  # 16x16 -> 8x8
+
+        # Fully connected layers
+        self.fc1 = nn.Linear(64 * 8 * 8, 128)
+        self.relu3 = nn.ReLU()
+        self.dropout = nn.Dropout(0.5)
+        self.fc2 = nn.Linear(128, num_classes)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Forward pass.
+
+        Args:
+            x: Input tensor of shape (batch_size, 3, 32, 32)
+
+        Returns:
+            torch.Tensor: Logits of shape (batch_size, num_classes)
+        """
+        x = self.pool1(self.relu1(self.bn1(self.conv1(x))))
+        x = self.pool2(self.relu2(self.bn2(self.conv2(x))))
+        x = x.view(x.size(0), -1)  # Flatten
+        x = self.dropout(self.relu3(self.fc1(x)))
+        x = self.fc2(x)
+        return x
+
+def load_model(model_path: str, device: str = "cuda") -> CIFAR10CNN:
+    """
+    Load trained CIFAR-10 model.
+
+    Args:
+        model_path: Path to model checkpoint (.pth file)
+        device: Device to load model on ('cuda' or 'cpu')
+
+    Returns:
+        CIFAR10CNN: Loaded model in eval mode
+    """
+    model = CIFAR10CNN(num_classes=10)
+
+    # Load checkpoint
+    checkpoint = torch.load(model_path, map_location=device)
+
+    # Handle both direct state_dict and checkpoint dict formats
+    if isinstance(checkpoint, dict) and "model_state_dict" in checkpoint:
+        model.load_state_dict(checkpoint["model_state_dict"])
+    else:
+        model.load_state_dict(checkpoint)
+
+    model = model.to(device)
+    model.eval()
+
+    return model
+
+def ifgsm_targeted_attack(
+    model: nn.Module,
+    image: torch.Tensor,
+    target_class: int,
+    epsilon: float,
+    mean: list,
+    std: list,
+    num_iterations: int = 50,
+    alpha: float = None,
+    device: str = "cuda",
+) -> torch.Tensor:
+    """
+    Implement Iterative FGSM targeted attack.
+
+    This implements the complete I-FGSM algorithm:
+    1. Start with clean image
+    2. For each iteration:
+       - Compute loss with respect to target class (minimize for targeted)
+       - Compute gradient of loss w.r.t. input
+       - Take small step in direction that decreases loss
+       - Project perturbation to L∞ ball
+       - Clip to valid range [0,1]
+
+    Args:
+        model: Trained model
+        image: Clean image tensor in [0,1] range, shape (3, 32, 32)
+        target_class: Target class to achieve
+        epsilon: L∞ perturbation budget
+        mean: Normalization mean per channel
+        std: Normalization std per channel
+        num_iterations: Number of iterations
+        alpha: Step size (defaults to epsilon/num_iterations)
+        device: Device to run attack on
+
+    Returns:
+        torch.Tensor: Adversarial image in [0,1] range
+    """
+    # Step size: divide epsilon by number of iterations for fine control
+    if alpha is None:
+        alpha = epsilon / num_iterations
+
+    # Convert normalization params to tensors
+    mean_t = torch.tensor(mean, device=device).view(3, 1, 1)
+    std_t = torch.tensor(std, device=device).view(3, 1, 1)
+
+    # Move to device and clone
+    x_adv = image.clone().to(device)
+    x_orig = image.clone().to(device)
+
+    # Target tensor
+    target = torch.tensor([target_class], device=device)
+
+    print(f"\n{'=' * 60}")
+    print(f"I-FGSM Targeted Attack")
+    print(f"{'=' * 60}")
+    print(f"Target class: {target_class}")
+    print(f"Epsilon: {epsilon:.6f} ({epsilon * 255:.1f}/255)")
+    print(f"Iterations: {num_iterations}")
+    print(f"Step size (alpha): {alpha:.6f}")
+    print(f"{'=' * 60}\n")
+
+    for iteration in range(num_iterations):
+        # Normalize image for model input
+        x_norm = (x_adv - mean_t) / std_t
+        x_norm.requires_grad = True
+
+        # Forward pass
+        outputs = model(x_norm.unsqueeze(0))
+
+        # Targeted attack: minimize loss toward target class
+        loss = F.cross_entropy(outputs, target)
+
+        # Backward pass to get gradients
+        model.zero_grad()
+        loss.backward()
+        grad_norm = x_norm.grad
+
+        # Convert gradient from normalized space to pixel space
+        # Chain rule: d/dx = d/dx_norm * dx_norm/dx = grad_norm * (1/std)
+        grad_pixel = grad_norm / std_t
+
+        # Targeted attack: move in NEGATIVE gradient direction (minimize loss)
+        x_adv = x_adv - alpha * grad_pixel.sign()
+
+        # Project to L∞ ball around original image
+        delta = x_adv - x_orig
+        delta = torch.clamp(delta, -epsilon, epsilon)
+        x_adv = x_orig + delta
+
+        # Clip to valid pixel range
+        x_adv = torch.clamp(x_adv, 0.0, 1.0)
+
+        # Detach for next iteration
+        x_adv = x_adv.detach()
+
+        if (iteration + 1) % 10 == 0:
+            # Check current prediction
+            with torch.no_grad():
+                x_norm_check = (x_adv - mean_t) / std_t
+                pred = model(x_norm_check.unsqueeze(0)).argmax(dim=1).item()
+            print(
+                f"  Iteration {iteration + 1}/{num_iterations} - Current prediction: {pred}"
+            )
+
+    print(f"\n{'=' * 60}")
+    print(f"Attack complete!")
+    print(f"{'=' * 60}\n")
+
+    return x_adv.detach().cpu()
+
+def test_api_robustness(host: str, image: torch.Tensor) -> bool:
+    """
+    Test API robustness against invalid inputs.
+
+    Args:
+        host: Challenge server host URL
+        image: Clean challenge image
+
+    Returns:
+        bool: True if all robustness tests pass
+    """
+    print(f"\n{'=' * 60}")
+    print("API Robustness Tests")
+    print(f"{'=' * 60}\n")
+
+    all_passed = True
+
+    # Test 1: Invalid JSON
+    print("[Test 1] Invalid JSON format...")
+    try:
+        response = requests.post(
+            f"{host}/submit",
+            data="invalid json",
+            headers={"content-type": "application/json"},
+        )
+        if response.status_code == 400:
+            print("  [+] Correctly rejected invalid JSON")
+        else:
+            print(f"  [-] Expected 400, got {response.status_code}")
+            all_passed = False
+    except Exception as e:
+        print(f"  [-] Exception: {e}")
+        all_passed = False
+
+    # Test 2: Missing image field
+    print("\n[Test 2] Missing 'image' field...")
+    try:
+        response = requests.post(f"{host}/submit", json={"wrong_field": "data"})
+        if response.status_code == 400:
+            result = response.json()
+            if "error" in result and "image" in result["error"].lower():
+                print("  [+] Correctly rejected missing image field")
+            else:
+                print(f"  [-] Expected error message about missing image field")
+                all_passed = False
+        else:
+            print(f"  [-] Expected 400, got {response.status_code}")
+            all_passed = False
+    except Exception as e:
+        print(f"  [-] Exception: {e}")
+        all_passed = False
+
+    # Test 3: Invalid base64
+    print("\n[Test 3] Invalid base64 encoding...")
+    try:
+        response = requests.post(
+            f"{host}/submit", json={"image": "not_valid_base64!!!"}
+        )
+        if response.status_code == 400:
+            print("  [+] Correctly rejected invalid base64")
+        else:
+            print(f"  [-] Expected 400, got {response.status_code}")
+            all_passed = False
+    except Exception as e:
+        print(f"  [-] Exception: {e}")
+        all_passed = False
+
+    # Test 4: Corrupted image data
+    print("\n[Test 4] Corrupted image data...")
+    try:
+        # Send corrupted base64 that decodes but isn't a valid PNG
+        corrupted_base64 = base64.b64encode(b"not a valid png image").decode("utf-8")
+        response = requests.post(f"{host}/submit", json={"image": corrupted_base64})
+        if response.status_code == 400:
+            print("  [+] Correctly rejected corrupted image data")
+        else:
+            print(f"  [-] Expected 400, got {response.status_code}")
+            all_passed = False
+    except Exception as e:
+        print(f"  [-] Exception: {e}")
+        all_passed = False
+
+    print(f"\n{'=' * 60}")
+    print(f"API Robustness: {'PASSED' if all_passed else 'FAILED'}")
+    print(f"{'=' * 60}\n")
+
+    return all_passed
+
+def test_anti_cheat(
+    host: str,
+    model: nn.Module,
+    image: torch.Tensor,
+    target_class: int,
+    epsilon: float,
+    mean: list,
+    std: list,
+    device: str,
+) -> bool:
+    """
+    Test anti-cheat mechanisms.
+
+    Args:
+        host: Challenge server host URL
+        model: Trained model
+        image: Clean challenge image
+        target_class: Target class for attack
+        epsilon: L∞ constraint
+        mean: Normalization mean
+        std: Normalization std
+        device: Device to run on
+
+    Returns:
+        bool: True if all anti-cheat tests pass
+    """
+    print(f"\n{'=' * 60}")
+    print("Anti-Cheat Mechanism Tests")
+    print(f"{'=' * 60}\n")
+
+    all_passed = True
+    mean_t = torch.tensor(mean).view(3, 1, 1)
+    std_t = torch.tensor(std).view(3, 1, 1)
+
+    # Test 1: Clean image (no modification)
+    print("[Test 1] Submitting clean image (should fail)...")
+    try:
+        clean_base64 = tensor_to_base64(image)
+        response = requests.post(f"{host}/submit", json={"image": clean_base64})
+        result = response.json()
+        if not result["success"] and not result["validation"]["target_achieved"]:
+            print("  [+] Correctly rejected clean image")
+        else:
+            print("  [-] Should reject clean image (no target achieved)")
+            all_passed = False
+    except Exception as e:
+        print(f"  [-] Exception: {e}")
+        all_passed = False
+
+    # Test 2: Perturbation too large
+    print("\n[Test 2] Perturbation exceeding epsilon (should fail)...")
+    try:
+        large_pert = image.clone()
+        large_pert = torch.clamp(large_pert + epsilon * 2, 0, 1)
+        large_base64 = tensor_to_base64(large_pert)
+        response = requests.post(f"{host}/submit", json={"image": large_base64})
+        result = response.json()
+        if not result["success"] and not result["validation"]["linf_satisfied"]:
+            print("  [+] Correctly rejected oversized perturbation")
+        else:
+            print("  [-] Should reject perturbation > epsilon")
+            all_passed = False
+    except Exception as e:
+        print(f"  [-] Exception: {e}")
+        all_passed = False
+
+    # Test 3: Patch-only attack (localized perturbation)
+    print("\n[Test 3] Localized patch attack (should fail)...")
+    try:
+        patch_image = image.clone()
+        # Only modify top-left 8x8 patch
+        patch_image[:, :8, :8] = torch.clamp(
+            patch_image[:, :8, :8] + epsilon * 0.9, 0, 1
+        )
+        patch_base64 = tensor_to_base64(patch_image)
+        response = requests.post(f"{host}/submit", json={"image": patch_base64})
+        result = response.json()
+        if not result["success"] and not result["validation"].get(
+            "sufficient_coverage", True
+        ):
+            print("  [+] Correctly rejected localized patch")
+        else:
+            print("  [-] Should reject patch-only attacks")
+            all_passed = False
+    except Exception as e:
+        print(f"  [-] Exception: {e}")
+        all_passed = False
+
+    # Test 4: Wrong gradient direction (reversed targeted attack)
+    print("\n[Test 4] Reversed gradient direction (should fail)...")
+    try:
+        x_wrong = image.clone().to(device)
+
+        # Do targeted attack but with WRONG sign (positive instead of negative)
+        for _ in range(30):
+            x_norm = (x_wrong - mean_t.to(device)) / std_t.to(device)
+            x_norm.requires_grad = True
+
+            outputs = model(x_norm.unsqueeze(0))
+            loss = F.cross_entropy(outputs, torch.tensor([target_class], device=device))
+            model.zero_grad()
+            loss.backward()
+
+            grad_norm = x_norm.grad
+            grad_pixel = grad_norm / std_t.to(device)
+
+            # WRONG: Use positive sign (should be negative for targeted)
+            x_wrong = x_wrong + (epsilon / 30) * grad_pixel.sign()
+            x_wrong = torch.clamp(x_wrong, 0, 1)
+
+        x_wrong = x_wrong.cpu()
+        wrong_base64 = tensor_to_base64(x_wrong)
+        response = requests.post(f"{host}/submit", json={"image": wrong_base64})
+        result = response.json()
+
+        # Should fail either due to not achieving target OR failing direction check
+        if not result["success"]:
+            if result["validation"].get("target_achieved") and not result[
+                "validation"
+            ].get("direction_check", True):
+                print("  [+] Correctly rejected wrong gradient direction")
+            elif not result["validation"].get("target_achieved"):
+                print("  [+] Wrong direction prevented target achievement")
+            else:
+                print("  [-] Should reject reversed gradients")
+                all_passed = False
+        else:
+            print("  [-] Should reject attacks with wrong gradient direction")
+            all_passed = False
+    except Exception as e:
+        print(f"  [-] Exception: {e}")
+        all_passed = False
+
+    print(f"\n{'=' * 60}")
+    print(f"Anti-Cheat Tests: {'PASSED' if all_passed else 'FAILED'}")
+    print(f"{'=' * 60}\n")
+
+    return all_passed
+
+def solve_challenge(host: str, device: str = "cuda", run_tests: bool = False):
+    """
+    Main solver function.
+
+    Args:
+        host: Challenge server host URL
+        device: Device for attack ('cuda' or 'cpu')
+        run_tests: Whether to run robustness and anti-cheat tests
+    """
+    print(f"\n{'=' * 60}")
+    print(f"Skills Assessment 1")
+    print(f"{'=' * 60}\n")
+
+    # Determine model path (assume running from solver/ directory)
+    solver_dir = os.path.dirname(os.path.abspath(__file__))
+    challenge_dir = os.path.dirname(solver_dir)
+    assets_dir = os.path.join(challenge_dir, "assets")
+    model_path = os.path.join(assets_dir, "cifar10_model_best.pth")
+
+    # Ensure assets directory exists
+    os.makedirs(assets_dir, exist_ok=True)
+
+    # Fetch model if missing
+    if not os.path.exists(model_path):
+        print(f"[1/5] Model not found at {model_path}")
+        weights_url = f"{host}/model/weights"  # Adjust if your endpoint differs
+        print(f"Attempting to download weights from: {weights_url}")
+
+        try:
+            resp = requests.get(weights_url, timeout=30)
+            resp.raise_for_status()
+            with open(model_path, "wb") as f:
+                f.write(resp.content)
+            print(f"  Weights downloaded and saved to {model_path}")
+        except requests.exceptions.RequestException as e:
+            print(f"[-] Error: Failed to download model weights from {weights_url}")
+            print(f"Details: {e}")
+            print("Hint: Ensure the challenge server exposes /model/weights and is reachable.")
+            raise
+
+    # Step 1: Load model
+    print("[1/5] Loading CIFAR-10 model...")
+    model = load_model(model_path, device=device)
+    print(f"  Model loaded on {device}")
+
+    # Step 2: Get challenge parameters
+    print("\n[2/5] Fetching challenge...")
+    response = requests.get(f"{host}/challenge")
+    response.raise_for_status()
+    challenge = response.json()
+
+    print(
+        f"  Original class: {challenge['original_class']} ({challenge['original_class_name']})"
+    )
+    print(
+        f"  Target class: {challenge['target_class']} ({challenge['target_class_name']})"
+    )
+    print(f"  Epsilon (L∞): {challenge['epsilon']:.6f}")
+    print(f"  Max iterations hint: {challenge['max_iterations_hint']}")
+
+    # Step 3: Decode challenge image
+    print("\n[3/5] Decoding challenge image...")
+    image = base64_to_tensor(challenge["image"])
+    mean = challenge["normalization"]["mean"]
+    std = challenge["normalization"]["std"]
+    print(f"  Image shape: {image.shape}")
+    print(f"  Normalization: mean={mean}, std={std}")
+
+    # Verify original prediction
+    with torch.no_grad():
+        mean_t = torch.tensor(mean).view(3, 1, 1)
+        std_t = torch.tensor(std).view(3, 1, 1)
+        img_norm = (image - mean_t) / std_t
+        orig_pred = model(img_norm.unsqueeze(0).to(device)).argmax(dim=1).item()
+    print(f"  Verified original prediction: {orig_pred}")
+
+    # Optional: Run robustness and anti-cheat tests
+    if run_tests:
+        print("\n" + "=" * 60)
+        print("RUNNING SECURITY TESTS")
+        print("=" * 60)
+
+        robustness_passed = test_api_robustness(host, image)
+        anticheat_passed = test_anti_cheat(
+            host,
+            model,
+            image,
+            challenge["target_class"],
+            challenge["epsilon"],
+            mean,
+            std,
+            device,
+        )
+
+        if robustness_passed and anticheat_passed:
+            print(f"\n{'=' * 60}")
+            print("[+] ALL SECURITY TESTS PASSED")
+            print(f"{'=' * 60}\n")
+        else:
+            print(f"\n{'=' * 60}")
+            print("[-] SOME SECURITY TESTS FAILED")
+            print(f"{'=' * 60}\n")
+
+        print("Proceeding with legitimate attack...\n")
+
+    # Step 4: Run I-FGSM attack
+    print("\n[4/5] Running I-FGSM targeted attack...")
+    adv_image = ifgsm_targeted_attack(
+        model=model,
+        image=image,
+        target_class=challenge["target_class"],
+        epsilon=challenge["epsilon"],
+        mean=mean,
+        std=std,
+        num_iterations=50,
+        device=device,
+    )
+
+    # Verify adversarial prediction
+    with torch.no_grad():
+        adv_norm = (adv_image - mean_t) / std_t
+        adv_pred = model(adv_norm.unsqueeze(0).to(device)).argmax(dim=1).item()
+    print(f"  Final adversarial prediction: {adv_pred}")
+
+    # Step 5: Submit adversarial example
+    print("\n[5/5] Submitting adversarial example...")
+    adv_base64 = tensor_to_base64(adv_image)
+
+    response = requests.post(f"{host}/submit", json={"image": adv_base64})
+    response.raise_for_status()
+    result = response.json()
+
+    # Display results
+    print(f"\n{'=' * 60}")
+    print("Validation Results")
+    print(f"{'=' * 60}")
+
+    validation = result["validation"]
+    print(f"L∞ norm: {validation['linf_norm']:.6f} / {challenge['epsilon']:.6f}")
+    print(f"L∞ constraint satisfied: {validation['linf_satisfied']}")
+    print(f"Valid range [0,1]: {validation['valid_range']}")
+    print(f"Adversarial prediction: {validation['adversarial_class']}")
+    print(f"Target class: {validation['target_class']}")
+    print(f"Target achieved: {validation['target_achieved']}")
+    print(f"\nSuccess: {result['success']}")
+
+    if result["success"]:
+        print(f"\n{'=' * 60}")
+        print(f"[+] FLAG: {result['flag']}")
+        print(f"{'=' * 60}\n")
+    else:
+        print(f"\n[-] {result['message']}")
+        if "hint" in result:
+            print(f"[!] {result['hint']}\n")
+
+def main():
+    """Main entry point."""
+    parser = argparse.ArgumentParser(
+        description="Working solver for Skills Assessment 1"
+    )
+    parser.add_argument(
+        "--host",
+        type=str,
+        default="http://154.57.164.76:31384",
+        help="Challenge server host URL (default: http://localhost:8000)",
+    )
+    parser.add_argument(
+        "--device",
+        type=str,
+        default="cuda" if torch.cuda.is_available() else "cpu",
+        help="Device to run attack on (default: cuda if available, else cpu)",
+    )
+    parser.add_argument(
+        "--skip-tests",
+        action="store_true",
+        help="Skip API robustness and anti-cheat tests (tests run by default)",
+    )
+    args = parser.parse_args()
+
+    try:
+        solve_challenge(args.host, args.device, not args.skip_tests)
+    except requests.exceptions.ConnectionError:
+        print(f"\n[-] Error: Could not connect to {args.host}")
+    except Exception as e:
+        print(f"\n[-] Error: {str(e)}\n")
+        raise
+
+
+if __name__ == "__main__":
+    main()