Alice's warm up —— pytorch / 搜索

题目描述

“Welcome to the XCTF-*CTF2022, I’m Alice and interested in AI security. I prepared a easy warm-up for you before you enjoy those pure AI security challenges.
Like humans, AI also needs to warm up before running. Can you find something strange in these initialized parameters?”

思路

• 关注模型的初始参数
• （不是）纯粹 ai 问题

第一步：加载模型

• 需要注意：AI 模型文件和当前的 AI 模型不同
  • AI 模型文件：完整清晰的模型结构的文件
  • 当前 AI 模型：更多的是黑盒
• 获取的文件为 xxx.zip，注意 pytorch 模型保存的文件结构，模型本身就是zip
  • 其结构就是：archive 文件夹，诺干个 AI 模型文件
• 因此，直接用torch.load导入zip文件即可

path = 'Alice_warm_up.zip'  
net = torch.load(path)

直接加载的时候，会报了一个错：

AttributeError: Can't get attribute 'AliceNet1' on <module '__main__' from 'xxxx/xxx.py'>

这是因为载入模型的时候，还需要先恢复模型，这里报错缺少AliceNet1类，那我们就直接补上一个空的类就可以

class AliceNet1(nn.Module):  
    pass

找到模型中的特殊点

• 看看网络结构：47x10x1的全连接网络

AliceNet1(  
  (fc): Sequential(    
	  (0): Linear(in_features=47, out_features=47, bias=True)    
	  (1): Linear(in_features=47, out_features=10, bias=True)    
	  (2): Linear(in_features=10, out_features=1, bias=True)  
  )
)  
47*10*1

• 输出 net 中的参数看看怎个事：发现除了第0层只有0和1，其他都是正常矩阵[-1,1]
  • 这里使用 pytorch 的 state_dict() 可以获取每层参数名

for name in net.state_dict(): 
	print(net.state_dict()[name])

fc.0.weight
tensor([[0., 0., 1.,  ..., 0., 0., 0.],
        [0., 0., 0.,  ..., 0., 0., 0.],
        [0., 0., 0.,  ..., 0., 0., 0.],
        ...,
        [0., 0., 0.,  ..., 0., 0., 0.],
        [0., 0., 0.,  ..., 0., 0., 0.],
        [0., 0., 0.,  ..., 0., 0., 0.]])
fc.0.bias
tensor([-0.0481,  0.0083,  0.0302,  0.0301, -0.0125,  0.1012,  0.0500,  0.0477,
         0.0310,  0.0276, -0.0967,  0.0889, -0.0678, -0.1428,  0.1141, -0.1254,
         0.1439, -0.0778,  0.0316, -0.0729,  0.0046, -0.1126, -0.0916,  0.1256,
         0.0641, -0.0317,  0.0422, -0.1378, -0.0489,  0.0612,  0.0500,  0.0190,
        -0.0649, -0.1008, -0.1050, -0.0399,  0.0906,  0.0275, -0.0825, -0.1431,
         0.1126, -0.0183,  0.1168, -0.0400, -0.0637,  0.0806, -0.0296])
fc.1.weight
tensor([[-0.0512,  0.1195, -0.0837, -0.0519, -0.0278,  0.0530,  0.1385, -0.0872,
         -0.1022, -0.0910, -0.0098,  0.0995,  0.0531, -0.0997, -0.0123, -0.0347,
         -0.0872,  0.0987, -0.0472,  0.0851, -0.1073, -0.0153, -0.0942,  0.0949,
          0.0522,  0.0521,  0.1063,  0.1335,  0.0305,  0.1082,  0.0114, -0.1429,
         -0.1264,  0.1127, -0.1318,  0.0350,  0.1166,  0.1224,  0.0600, -0.0837,
         -0.0425, -0.0854,  0.0214, -0.1391, -0.0359, -0.0529, -0.0379],
         ...
        [-0.0641, -0.1137, -0.0556,  0.1383, -0.0967,  0.0524, -0.0661,  0.0510,
         -0.1030, -0.0732,  0.1109,  0.1101, -0.1164, -0.0505, -0.0610, -0.0219,
         -0.1451, -0.0486, -0.0898, -0.1229,  0.1050, -0.0934, -0.0408,  0.0432,
          0.0159,  0.0220,  0.0875, -0.0512, -0.0437,  0.0833,  0.0277,  0.0892,
         -0.1136, -0.1330, -0.0778,  0.0363,  0.0043,  0.1038, -0.1079, -0.0030,
         -0.0358, -0.1302,  0.0822,  0.0155,  0.0218,  0.0417,  0.1241]])
fc.1.bias
tensor([ 0.0427,  0.1376, -0.0805,  0.1418,  0.1263,  0.0791, -0.0008,  0.0997,
         0.0499, -0.0104])
fc.2.weight
tensor([[-0.2986, -0.2412,  0.2308, -0.0410,  0.1583, -0.0777, -0.2521,  0.0848,
         -0.0169, -0.0387]])
fc.2.bias
tensor([-0.1993])

• 再看看 hint.py：
  • flag 长度为16
  • flag 内容范围以及格式：可以通过print(flagset)查看0123456789abcdefghijklmnopqrstuvwxyz*CTF{ALIZE}
• 同时发现，这里 flagset 中的字母数量正好是47
• 因此可以大胆猜测这里的字母和第0层的01矩阵行列是相对应的
• => 这个矩阵是一个邻接矩阵

通过输出这个矩阵图，观察

• 进一步的分析会发现，整个邻接矩阵代表的是一个带环的有向图。—— 是字母之间的联系

因此，想到了使用

• 使用 dfs 算法，最大深度16(flag长度)，初始位置"*"

完整exp：

import torch  
import torch.nn as nn  
import string  
import matplotlib.pyplot as plt  
import numpy as np  
  
  
class AliceNet1(nn.Module):  
    pass  
  
  
  
  
  
'''  
AliceNet1(  
  (fc): Sequential(    (0): Linear(in_features=47, out_features=47, bias=True)    (1): Linear(in_features=47, out_features=10, bias=True)    (2): Linear(in_features=10, out_features=1, bias=True)  ))  
47*10*1  
  
(36,37)->"*" 
(37,38)->"C"  
(38,39)->"T"  
(39,40)->"F"  
(40,41)->"{"  
  
'''  
  
# for name in net.state_dict():  
# print(name)  
# print(net.state_dict()[name])  
  
  
# print(mymat[36][37])  
  
# 绘制第一层参数图  
# plt.imshow(mymat, cmap='binary', interpolation='nearest')  
# plt.title('First Layer Parameters')  
# plt.colorbar()  
# plt.xlabel('Input Size')  
# plt.ylabel('Hidden Size')  
# plt.xticks(np.arange(0, mymat.shape[1], 1))  
# plt.yticks(np.arange(0, mymat.shape[0], 1))  
# plt.grid(True)  
# plt.show()  
  
# print(string.printable[0:36].find('r'))  
  
def char2num(ch):  
    tmpset = string.printable[0:36] + '*CTF{ALIZE}'  
    tmplen = len(tmpset)  
    # print(tmpset)  
    for i in range(tmplen):  
        if (ch == tmpset[i]):  
            return i  
  
def dfs(ch,depth,ans):  
    ans += ch  
    if len(ans)==flaglen and ans[-1]=='}':  
        print('get flag: ',ans)  
        exit(0)  
    elif len(ans)==flaglen:  
        return  
    else:  
  
        tmpi = char2num(ch)  
        for i in range(setlen):  
            if flagset[i]==ch:  
                continue  
            tmpj = char2num(flagset[i])  
            if mymat[tmpi][tmpj]==1.0 and used[tmpj]==False:  
                # print(ans)  
                used[tmpj]=True  
                dfs(flagset[i],depth+1,ans)  
                used[tmpj]=False  
  
  
path = 'Alice_warm_up.zip'  
net = torch.load(path)  
# print(net)  
mymat = net.state_dict()['fc.0.weight'].tolist()  
flagset = string.printable[0:36] + '*CTF{ALIZE}'  
print(flagset)  
setlen = len(flagset)  
flaglen = 16  
used = [0] * setlen  
flag = ''  
used[char2num('*')]=1  
dfs('*',0,flag)

得到flag：*CTF{qx1jukznmr}