今年强网杯出了一道关于神经网络的Pwn题,当时看到感觉一脸懵逼。赛后看了writeup感觉也是一知半解。感觉个人在算法和模型上的知识欠缺较多,遂借这个机会学习下神经网络的相关知识。
神经网络是一种模拟人脑神经元工作原理的计算模型,它由大量的人工神经元相互连接构成,能够学习和处理各种复杂的数据模式。
加权求和阶段:
net_input = Σ(input_i * weight_i) + bias
其中,input_i
是第 i 个输入信号,weight_i
是对应输入信号的权重,bias
是一个偏置项,用于调整神经元的激活门槛。
激活函数阶段:
output = activation_function(net_input)
其中,activation_function
是所选的非线性函数,常见的激活函数包括sigmoid、ReLU(Rectified Linear Unit)、tanh、softmax等。
神经网络通过这种神经元的计算过程以及前向传播和反向传播的学习机制,能够在各种任务中表现出强大的学习和泛化能力,包括图像识别、语音识别、自然语言处理、推荐系统等领域。
卷积神经网络(Convolutional Neural Networks, CNN)是一种专门设计用于处理图像、视频、语音等具有网格结构数据的深度学习模型。
卷积层:这是CNN的核心层,使用一组可学习的过滤器(或称卷积核)在输入数据上进行滑动窗口操作,并执行元素级别的乘法和加法运算,生成特征图(或称为激活图)。这些特征图捕获输入中的局部特征和模式。
激活函数:如ReLU(Rectified Linear Unit)、sigmoid或tanh,用于引入非线性特性,使得网络能够学习更复杂的模式。
池化层:通常在卷积层之后,池化层通过取每个子区域的最大值、平均值或其他统计量来降低数据维度和减少计算复杂性,同时保留重要的特征信息。
全连接层:经过一系列的卷积和池化层后,可能会添加一个或多个全连接层。这些层中,每个神经元与前一层的所有神经元都有连接,用于将学到的局部特征组合成全局表示,并用于分类或回归任务。
参数共享:在同一特征图中,卷积核在所有位置重复使用,减少了模型的参数数量,提高了计算效率并增强了模型的泛化能力。
平移不变性:由于卷积操作对输入数据的位置相对不敏感,CNN能够在一定程度上保持对对象位置变化的不变性。
层次化表示:随着网络深度的增加,CNN能够学习越来越抽象和复杂的特征,从边缘和纹理到高级物体部分和整体形状。
卷积神经网络因其在处理图像和其他网格结构数据方面的优秀性能和高效性而被广泛应用于各种计算机视觉和深度学习任务中。
卷积神经网络(CNN)的计算过程可以分为以下几个主要步骤:
通过以上步骤,卷积神经网络能够自动从输入数据中学习和提取有用的特征,并用于各种任务,如图像分类、物体检测、语义分割和语音识别等。
第一遍分析题目,可以发现题目的验证流程大概如下:
以上部分是第一遍分析题目时可以得到的最基本的信息。再加上题目提示 "hint for you, norm, 0.1307, 0.3081" , 检索0.1307,0.3081这两个关键浮点数会发现这两个浮点数是MNIST数据集标准化的常用参数。
得到以上信息之后,结合题目名 EZNN(Neural Networks) 可以看出,题目的前半部分应该是要求输入一个神经网络的一些参数,然后以这些参数构建模型去运行对应的测试集。问题的关键在于如何正确的对程序进行输入,这里就需要对程序进行逆向。
根据字符串的一些提示,我们可以大概推断以下几个函数的名称或作用:
这里我们可以先尝试用pytroch提供的针对MINIST的数据集的模型样例跑一下,尝试训练我们第一个卷积神经网络的模型(安装完几个依赖库可以直接跑起来):
得到第一个模型之后我们需要确认,我们如何将这个模型输入给EZNN这个程序。这时候可以看一下pytorch提供的导出模型的接口,具体导出了哪些数据。导出一个CNN模型的样例:
torch.save(model.state_dict(), "mnist_cnn.pt")
可以看到主要保存的是模型的 state_dict(), state_dict() 里有什么呢?我们可以打印看一下:
import torch
state_dict = torch.load('mnist_cnn.pt')
print(state_dict)
看一下对应的输出,你会发现它包含了conv1.weight
,conv1.bias
,conv2.bias
,fc1.weight
,fc1.bias
,fc2.weight
,fc2.bias
。参考上面卷积神经网络的计算过程以及模型样例的代码,可以确认,这些就是最终训练得到的模型关键参数。
这时候可以对比一下模型样例中forward
的实现:
def forward(self, x):
x = self.conv1(x) #进行第一层卷积层的计算
x = F.relu(x) #激活函数
x = self.conv2(x) #进行第二层卷积层的计算
x = F.relu(x) #激活函数
x = F.max_pool2d(x, 2) #池化,kernel size是2
x = self.dropout1(x) #丢弃
x = torch.flatten(x, 1) #在进行连接层之前进行扁平化
x = self.fc1(x) #连接层第一次计算
x = F.relu(x) #激活函数
x = self.dropout2(x) #丢弃
x = self.fc2(x) #连接层第二次计算
output = F.log_softmax(x, dim=1)#得到最后的输出
return output
和题目给出的sub_11ECC0(运行模型)
的函数进行一个简单对比。可以粗略看出两者十分类似,都是对输入进行一系列的卷积,池化操作,最后进行全连接层(FC Layer)的计算。题目的输入前读取的两个整形变量最终分别会决定在sub_11ECC0(运行模型)
中进行卷积计算的次数和全连接层计算的次数。这里大致可以猜出,我们的输入应该就是pytorch导出的state_dict中的所有数据了。
直接发送模型样例代码训练出来的模型会发现输出“invalid!”。这时候需要对比下第一部分卷积层计算的部分样例代码和题目环境的不同,会发现卷积层使用的卷积核大小及参数不同:
def __init__(self):
super(Net, self).__init__()
self.conv1 = nn.Conv2d(1, 32, 3, 1)#参数分别为:in_channels,out_channels,kernel_size,stride
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)
而题目中生成卷积核的函数为sub_123F60
,通过逆向及参数分析可以看出题目创建的所有卷积核参数均为:输入通道为1,输出通道为1,3*3大小。结合程序逆向,我们修改样例模型为:
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = nn.Conv2d(1, 1, 3, 1)
self.conv2 = nn.Conv2d(1, 1, 3, 1)
self.dropout1 = nn.Dropout(0.25)
self.dropout2 = nn.Dropout(0.5)
self.fc1 = nn.Linear(121, 121)
self.fc2 = nn.Linear(121, 10)
def forward(self, x):
x = self.conv1(x)
x = F.relu(x)
x = F.max_pool2d(x, 2)
x = self.dropout1(x)
x = self.conv2(x)
x = F.relu(x)
x = torch.flatten(x, 1)
x = self.fc1(x)
x = F.relu(x)
x = self.dropout2(x)
x = self.fc2(x)
x = F.relu(x)
output = F.log_softmax(x, dim=1)
return output
以该模型进行输入之后,可以看到已经可以正常过第一次模型的check了:
这是可以发现题目输出了十个bacodoor的数据,标签,及其预测的标签。之后又是一组类似于第一步的输入。很明显,题目要求我们输入一个新模型,该新模型在题目给出的十个backdoor图片上应该符合题目要求的pridict label(即将图片识别为一个错误的结果)。同时我们通过逆向分析可以看到最后还会使用我们输入的第二个模型去跑目录在./testset/
下的数据集。至此我们可以总结一下题目的大致流程:
至此,这个题目的目的就是考察选手训练一个能识别出后门数据的模型。因为在backdoor数据集之后还会测试正常数据,这里我们可以考虑直接增大backdoor数据集在总训练集中的比例(或用backdoor多训几轮),来让该模型可以记住对应的backdoor。
剩下的一些工作就集中在数据处理等部分了,这里就不详细描述了,exp最终完成如下:
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
from torch.optim.lr_scheduler import StepLR
import hashlib
from pwn import *
#context.log_level = 'debug'
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = nn.Conv2d(1, 1, 3, 1)
self.conv2 = nn.Conv2d(1, 1, 3, 1)
self.dropout1 = nn.Dropout(0.25)
self.dropout2 = nn.Dropout(0.5)
self.fc1 = nn.Linear(121, 121)
self.fc2 = nn.Linear(121, 10)
def forward(self, x):
x = self.conv1(x)
x = F.relu(x)
x = F.max_pool2d(x, 2)
x = self.dropout1(x)
x = self.conv2(x)
x = F.relu(x)
x = torch.flatten(x, 1)
x = self.fc1(x)
x = F.relu(x)
x = self.dropout2(x)
x = self.fc2(x)
x = F.relu(x)
output = F.log_softmax(x, dim=1)
return output
def train(args, model, device, train_loader, optimizer, epoch):
model.train()
for batch_idx, (data, target) in enumerate(train_loader):
data, target = data.to(device), target.to(device)
optimizer.zero_grad()
output = model(data)
loss = F.nll_loss(output, target)
loss.backward()
optimizer.step()
if batch_idx % 10 == 0:
print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
epoch, batch_idx * len(data), len(train_loader.dataset),
100. * batch_idx / len(train_loader), loss.item()))
def test(model, device, test_loader):
model.eval()
test_loss = 0
correct = 0
with torch.no_grad():
for data, target in test_loader:
data, target = data.to(device), target.to(device)
output = model(data)
test_loss += F.nll_loss(output, target, reduction='sum').item() # sum up batch loss
pred = output.argmax(dim=1, keepdim=True) # get the index of the max log-probability
correct += pred.eq(target.view_as(pred)).sum().item()
test_loss /= len(test_loader.dataset)
print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
test_loss, correct, len(test_loader.dataset),
100. * correct / len(test_loader.dataset)))
#用backdoor的数据结合minst数据集进行训练
def train_with_backdoor(data):
torch.manual_seed(1)
device = torch.device("cpu")
train_kwargs = {'batch_size': 64,'shuffle': True}
test_kwargs = {'batch_size': 1000}
transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))
])
dataset1 = datasets.MNIST('../data', train=True, download=True,
transform=transform)
dataset2 = datasets.MNIST('../data', train=False,
transform=transform)
mixed = torch.utils.data.ConcatDataset([dataset1, data * 2000])
train_loader = torch.utils.data.DataLoader(mixed,**train_kwargs)
test_loader = torch.utils.data.DataLoader(dataset2, **test_kwargs)
backdoor_loader = torch.utils.data.DataLoader(data, **test_kwargs)
model = Net().to(device)
optimizer = optim.Adadelta(model.parameters(), lr=1.0)
#scheduler = StepLR(optimizer, step_size=1, gamma=0.7)
for epoch in range(1, 4):
train(args, model, device, train_loader, optimizer, epoch)
test(model, device, backdoor_loader)
test(model, device, test_loader)
#scheduler.step()
return model
def find_target(src,target):
for i in range(255):
for j in range(255):
for k in range(255):
dst = src.decode() + ('%02x%02x%02x'%(i,j,k))
t = hashlib.sha256()
t.update(bytes.fromhex(dst))
res = t.hexdigest()
if res == target.decode():
return dst
#s = remote('47.100.169.26',20232)
state_dict = torch.load("mnist_cnn.pt")
#s = process("./NN")
s = remote("127.0.0.1",1234)
s.recvuntil(b'buf (10 bytes) prefix:')
src = s.recvuntil(b'\n')[:-1].replace(b' ',b'')
print(src)
s.recvuntil(b'sha256(buf) = ')
target = s.recvuntil(b'\n')[:-1]
print(target)
s.recvuntil(b'please input the left of this buf (hex, ei.g., 56 12 8f):')
answer = find_target(src,target)[-6:]
print('Answer: '+answer)
s.sendline(' '+answer[:2]+' '+answer[2:4]+' '+answer[4:])
n_conv = 2
s.sendlineafter(b"> ", str(n_conv).encode())
for i in range(1, n_conv+1):
# weight
s.sendafter(b"$ ", state_dict[f"conv{i}.weight"].numpy().tobytes())
# bias
s.sendafter(b"$ ", state_dict[f"conv{i}.bias"].numpy().tobytes())
#gdb.attach(s)
n_fc = 2
# FC
s.sendlineafter(b"> ", str(n_fc).encode())
for i in range(1, n_fc+1):
# weight
s.sendafter(b"$ ", state_dict[f"fc{i}.weight"].numpy().tobytes())
# bias
s.sendafter(b"$ ", state_dict[f"fc{i}.bias"].numpy().tobytes())
#Now try to get backdoor data and train it.
train_data = []
s.recvuntil(b'Now you need to create a backdoor in your model!')
device = torch.device("cpu")
for i in range(10):
s.recvuntil(b"backdoor ")
s.recvuntil(b"\n")
data = [[0 for i in range(0,28)] for j in range(0,28)] # 创建数组接收数据
for i in range(0,28):
t = s.recvuntil(b'\n')[:-1].replace(b' ',b' ')
t = t.split(b' ')[:-1]
#print(t)
for j in range(0,28):
data[i][j] = float(t[j])
b = []
b.append(data)
backdoor_ts = torch.tensor(b,dtype=torch.float32).to(device)
s.recvuntil(b'predict label -> ')
label = s.recvuntil(b'\n')[:-1]
train_data.append((backdoor_ts, int(label)))
model = train_with_backdoor(train_data)
state_dict = model.state_dict()
n_conv = 2
for i in range(1, n_conv+1):
# weight
s.sendafter(b"$ ", state_dict[f"conv{i}.weight"].numpy().tobytes())
# bias
s.sendafter(b"$ ", state_dict[f"conv{i}.bias"].numpy().tobytes())
n_fc = 2
# FC
for i in range(1, n_fc+1):
# weight
s.sendafter(b"$ ", state_dict[f"fc{i}.weight"].numpy().tobytes())
# bias
s.sendafter(b"$ ", state_dict[f"fc{i}.bias"].numpy().tobytes())
s.interactive()
用pytorch训练模型
有一定的经验,且有一定的逆向能力。