added material
This commit is contained in:
Jeremy Janella
@@ -0,0 +1,210 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import base64
|
||||
import io
|
||||
import json
|
||||
from dataclasses import dataclass
|
||||
import os
|
||||
import time
|
||||
from typing import Tuple
|
||||
import numpy as np
|
||||
import requests
|
||||
from PIL import Image
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
# Define MNIST normalization constants
|
||||
MNIST_MEAN = 0.1307 # average pixel intensity of MNIST images scaled to [0,1]
|
||||
MNIST_STD = 0.3081 # standard deviation of pixel intensities in [0,1]
|
||||
|
||||
class SimpleClassifier(nn.Module):
|
||||
"""CNN matching the server-side architecture with log-softmax outputs."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.conv1 = nn.Conv2d(1, 32, 3, 1)
|
||||
self.conv2 = nn.Conv2d(32, 64, 3, 1)
|
||||
self.dropout1 = nn.Dropout(0.25)
|
||||
self.dropout2 = nn.Dropout(0.5)
|
||||
self.fc1 = nn.Linear(9216, 128)
|
||||
self.fc2 = nn.Linear(128, 10)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
x = self.conv1(x)
|
||||
x = torch.relu(x)
|
||||
x = self.conv2(x)
|
||||
x = torch.relu(x)
|
||||
x = torch.max_pool2d(x, 2)
|
||||
x = self.dropout1(x)
|
||||
x = torch.flatten(x, 1)
|
||||
x = self.fc1(x)
|
||||
x = torch.relu(x)
|
||||
x = self.dropout2(x)
|
||||
x = self.fc2(x)
|
||||
return torch.log_softmax(x, dim=1)
|
||||
|
||||
|
||||
def mnist_normalize(x01: torch.Tensor) -> torch.Tensor:
|
||||
"""Normalize a [0,1] tensor to MNIST stats for the classifier."""
|
||||
return (x01 - MNIST_MEAN) / MNIST_STD
|
||||
|
||||
def png_from_x01(x01: np.ndarray) -> str:
|
||||
"""Encode a `[0,1]` grayscale image `(28,28)` to base64 PNG string."""
|
||||
x255 = np.clip((x01 * 255.0).round(), 0, 255).astype(np.uint8)
|
||||
img = Image.fromarray(x255, mode="L")
|
||||
buf = io.BytesIO()
|
||||
img.save(buf, format="PNG", optimize=True)
|
||||
return base64.b64encode(buf.getvalue()).decode("ascii")
|
||||
|
||||
|
||||
def x01_from_b64_png(b64: str) -> np.ndarray:
|
||||
"""Decode base64 PNG to `[0,1]` numpy array of shape `(28,28)`."""
|
||||
raw = base64.b64decode(b64)
|
||||
img = Image.open(io.BytesIO(raw)).convert("L")
|
||||
if img.size != (28, 28):
|
||||
raise ValueError("Expected 28x28 PNG")
|
||||
x = np.asarray(img, dtype=np.float32) / 255.0
|
||||
return np.clip(x, 0.0, 1.0)
|
||||
|
||||
@dataclass
|
||||
class Challenge:
|
||||
epsilon: float
|
||||
label: int
|
||||
sample_index: int
|
||||
x01: np.ndarray # (1,1,28,28)
|
||||
|
||||
def fetch_challenge(host: str, retries: int = 30, delay: float = 1.0) -> Challenge:
|
||||
"""Fetch challenge with simple retry/backoff to tolerate startup races."""
|
||||
last_err = None
|
||||
for _ in range(max(1, retries)):
|
||||
try:
|
||||
r = requests.get(f"{host}/challenge", timeout=5)
|
||||
r.raise_for_status()
|
||||
payload = r.json()
|
||||
break
|
||||
except Exception as e: # noqa: BLE001
|
||||
last_err = e
|
||||
time.sleep(delay)
|
||||
else:
|
||||
raise RuntimeError(
|
||||
f"Failed to connect to {host}/challenge: {last_err}"
|
||||
) from last_err
|
||||
x2d = x01_from_b64_png(payload["image_b64"]) # (28,28)
|
||||
x4d = x2d[None, None, ...]
|
||||
return Challenge(
|
||||
epsilon=float(payload["epsilon"]),
|
||||
label=int(payload["label"]),
|
||||
sample_index=int(payload["sample_index"]),
|
||||
x01=x4d.astype(np.float32),
|
||||
)
|
||||
|
||||
def load_model(weights_path: str) -> SimpleClassifier:
|
||||
model = SimpleClassifier()
|
||||
state = torch.load(weights_path, map_location=torch.device("cpu"))
|
||||
model.load_state_dict(state)
|
||||
model.eval()
|
||||
return model
|
||||
|
||||
def fgsm_untargeted(
|
||||
model: nn.Module, x01: np.ndarray, y: int, epsilon: float
|
||||
) -> np.ndarray:
|
||||
"""Craft FGSM adversarial example under L_inf in [0,1].
|
||||
|
||||
Parameters
|
||||
----------
|
||||
model : nn.Module
|
||||
Classifier in eval mode.
|
||||
x01 : np.ndarray
|
||||
Baseline `[0,1]` image with shape `(1,1,28,28)`.
|
||||
y : int
|
||||
True class label.
|
||||
epsilon : float
|
||||
L_inf bound in `[0,1]` space.
|
||||
|
||||
Returns
|
||||
-------
|
||||
np.ndarray
|
||||
Adversarial image in `[0,1]` with shape `(1,1,28,28)`.
|
||||
"""
|
||||
x = torch.from_numpy(x01).float().requires_grad_(True)
|
||||
target = torch.tensor([y], dtype=torch.long)
|
||||
logits = model(mnist_normalize(x))
|
||||
loss = nn.NLLLoss()(logits, target)
|
||||
model.zero_grad()
|
||||
loss.backward()
|
||||
grad_sign = x.grad.detach().sign()
|
||||
x_adv = x + epsilon * grad_sign
|
||||
x_adv = torch.clamp(x_adv, 0.0, 1.0)
|
||||
return x_adv.detach().cpu().numpy()
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument(
|
||||
"--host", default="http://154.57.164.80:31923", help="Server base URL"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--weights", default="solver/fgsm_weights.pth", help="Path to model weights (downloads from /weights if missing)"
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
chall = fetch_challenge(args.host)
|
||||
if not os.path.exists(args.weights):
|
||||
os.makedirs(os.path.dirname(args.weights), exist_ok=True)
|
||||
wb = requests.get(f"{args.host}/weights", timeout=15).content
|
||||
with open(args.weights, "wb") as f:
|
||||
f.write(wb)
|
||||
model = load_model(args.weights)
|
||||
|
||||
# Local sanity prediction on clean image
|
||||
x = torch.from_numpy(chall.x01)
|
||||
clean_pred = int(torch.argmax(model(mnist_normalize(x)), dim=1).item())
|
||||
if clean_pred != chall.label:
|
||||
print(
|
||||
f"Warning: local clean prediction {clean_pred} != label {chall.label}; proceeding regardless"
|
||||
)
|
||||
|
||||
# Try conservative epsilons to account for 8-bit PNG quantization
|
||||
safety = [1.0 / 255.0 + 1e-4, 2e-3, 1e-3, 5e-4, 0.0]
|
||||
x_adv = None
|
||||
adv_pred = None
|
||||
linf = None
|
||||
for m in safety:
|
||||
eps_eff = max(0.0, chall.epsilon - m)
|
||||
cand = fgsm_untargeted(model, chall.x01, chall.label, eps_eff)
|
||||
cand_pred = int(
|
||||
torch.argmax(model(mnist_normalize(torch.from_numpy(cand))), dim=1).item()
|
||||
)
|
||||
if cand_pred != chall.label:
|
||||
x_adv = cand
|
||||
adv_pred = cand_pred
|
||||
linf = float(np.max(np.abs(x_adv - chall.x01)))
|
||||
break
|
||||
|
||||
if x_adv is None:
|
||||
# Fall back to slight under-shoot of epsilon
|
||||
eps_eff = chall.epsilon * 0.99
|
||||
x_adv = fgsm_untargeted(model, chall.x01, chall.label, eps_eff)
|
||||
adv_pred = int(
|
||||
torch.argmax(model(mnist_normalize(torch.from_numpy(x_adv))), dim=1).item()
|
||||
)
|
||||
linf = float(np.max(np.abs(x_adv - chall.x01)))
|
||||
|
||||
print(
|
||||
json.dumps({"linf": linf, "pred": adv_pred, "clean_pred": clean_pred}, indent=2)
|
||||
)
|
||||
|
||||
# Submit
|
||||
x2d = x_adv[0, 0]
|
||||
b64 = png_from_x01(x2d)
|
||||
r = requests.post(f"{args.host}/submit", json={"image_b64": b64}, timeout=10)
|
||||
try:
|
||||
r.raise_for_status()
|
||||
except Exception:
|
||||
print("Server response:", r.text)
|
||||
raise
|
||||
print("Flag:", r.json().get("flag"))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user