I wrote this as a companion challenge to attack of the worm for UMDCTF 2024. Except this time it's a patch attack.
once again, the fremen are trying to sabotage the spice harvest and need you to fool the spice harvester's
worm image recognition! this time though, you need to generate a 40x40 patch to add to the image.
We are provided model.pt storing model weights and server.py, but no image of a worm.
#!/usr/local/bin/python
import base64
import os
import sys
import numpy as np
from PIL import Image
import torch
import torch.nn as nn
from torchvision.models import resnet18
import torchvision.transforms as T
def mask_generation(patch=None, image_size=(3, 224, 224)):
applied_patch = np.zeros(image_size)
rotation_angle = np.random.choice(4)
for i in range(patch.shape[0]):
patch[i] = np.rot90(patch[i], rotation_angle)
x_location = np.random.randint(low=0, high=image_size[1]-patch.shape[1])
y_location = np.random.randint(low=0, high=image_size[2]-patch.shape[2])
applied_patch[:, x_location:x_location + patch.shape[1], y_location:y_location + patch.shape[2]] = patch
mask = applied_patch.copy()
mask[mask != 0] = 1.0
return applied_patch, mask, x_location, y_location
model = resnet18()
model.fc = nn.Linear(model.fc.in_features, 1)
model.load_state_dict(torch.load('model.pt'))
model.eval()
inp = input(f"Enter a base64 encoded 40x40 image patch.\n")
try:
patch = np.frombuffer(base64.standard_b64decode(inp), dtype=np.uint8)
patch = patch.reshape((40, 40, 3)).transpose(2, 0, 1) / 255
except:
print("Invalid patch.")
sys.exit(0)
total = 0
correct = 0
for test_img in os.listdir("worms"):
total += 1
img = Image.open(f"worms/{test_img}")
img = T.Resize(size=(224, 224))(img)
img = T.ToTensor()(img)
applied_patch, mask, _, _ = mask_generation(patch, image_size=(3, 224, 224))
applied_patch = torch.from_numpy(applied_patch)
mask = torch.from_numpy(mask)
perturbated_image = mask.float() * applied_patch.float() + (1 - mask.float()) * img.float()
output = model(perturbated_image.unsqueeze(0))
predicted = (output > 0).int()
if predicted == 0:
correct += 1
if correct / total >= 0.7:
with open("flag.txt", "r") as f:
print("LISAN AL GAIB")
print(f.read())
else:
print("Worm detected.")
The server takes your patch as input, then applies it to an unprovided worm image at a random location in the image, and with a random rotation. It then runs the classification model on this modified image and counts if it was correct. This occurs for every image in the folder of arbitrary worm images, and the flag is provided if at least 70% of the images are misclassified.
I think it's a very interesting problem - the same patch needs to lead to misclassification for multiple arbitrary worm images, and it also needs to be somewhat invariant to translation and rotation.
It's really an optimization problem at its core, and is actually a classic adversarial machine learning problem, stemming all the way back to 2017 with this paper: https://arxiv.org/abs/1712.09665.
My solve script is a little long and largely copied from patch attack implementations available online, but I'll describe the basic idea.
You essentially write a standard training loop, but instead optimizing your 40x40 patch image. You collect a bunch of images of worms from online to use as training data. Then at every training iteration, you evaluate the model on the modified image, compute a loss and backpropagate against the negative loss to attempt to misclassify, and optimize the patch.
It's honestly a very standard optimization solution, though the code ends up being a little lengthy due to having to load training data and deal with patch application.
Here's the patch we obtain:
And we get the flag: UMDCTF{sandworms_love_adversarial_patches}. This challenge only got 1 solve during the competition, which was a little surprising but I guess this hardcore ML is a little unconventional in CTFs.
My solve script along with the challenge source is available here: https://github.com/UMD-CSEC/UMDCTF-2024-Challenges/tree/main/misc/the-worm-strikes-back