2022 *CTF Misc babyFL WP（AI安全）

题目描述：

源码链接：

链接：https://pan.baidu.com/s/14eOkeEAxO5byopzZhgyh8A
提取码：pl5w

推荐使用jupyter

题目分析：

先看源码

if __name__ == '__main__':
    try:
        if not os.path.exists('/home/ctf/model'):
            os.mkdir("/home/ctf/model")
            train_models() #训练模型
        parameters = load_parameters() #加载正确识别的模型的参数
        parameters = get_input_parameter(parameters) #添加自己输入的权重
        model = aggregation(parameters) #用更新后的权重组合更新模型
        test(model) #测试模型，错误识别数字的正确率高于95%即可获取flag
    except Exception as e:
        traceback.print_exc()
        print("error")

分析可知，需要自己输入一组参数来使得acc>0.95来获取flag。

观察load_test_data函数可知，测试所用的数据是错误的，与正确识别的数字标签相反（0被错误识别为9），故整体思路为，分别训练正确识别和错误识别两个模型，然后用21*错误模型参数权重-n*正确模型参数权重=要输入的参数权重获取输入参数并保存。(n取0-4中的数即可)

题目解答：

训练正确模型：

def train_models_true():
    (x, y), (_, _) = mnist.load_data()
    x = x.reshape(-1, 28, 28, 1)
    model = new_model()
    model.fit(x, y, batch_size=64, epochs=30)
    model.save("1.h5")#正确模型的参数

训练错误模型：

def load_test_data():
    (x, y), (_, _) = mnist.load_data()
    l = len(y)
    for i in range(l):
        y[i] = 9 - y[i]
    x = x.reshape(-1, 28, 28, 1)
    return x, y

def train_models_false():
    x, y = load_test_data()
    x = x.reshape(-1, 28, 28, 1)
    model = new_model()
    model.fit(x, y, batch_size=64, epochs=30)
    model.save("2.h5")#错误模型参数

保存权重：

model1 = load_model('2.h5')
param2 = model1.get_weights()
model2 = load_model('1.h5')
param1 = model2.get_weights()
weights = []
ans = []
for i in range(len(param2)):
    f = open(str(i)+'.txt', 'a')
    layer1 = param1[i]
    layer2 = param2[i]
    sum = 0
    sum = layer2*21-layer1
    weights.append(sum)
    arr = np.array(sum, float)
    arr = arr.ravel()
    for i in range(len(arr)):
        f.write(str(arr[i])+' ')
    f.write('\n')
f.close()

搭建本地环境并测试：

import os
import traceback

import numpy as np
from tensorflow.keras import Sequential
from tensorflow.keras.layers import  Dense, Conv2D, Flatten, MaxPooling2D
from tensorflow import keras

from tensorflow.keras.models import load_model
from tensorflow.keras.datasets import mnist

participant_number = 20

def new_model():
    model = Sequential()
    model.add(Conv2D(10, (3, 3), input_shape=(28, 28, 1)))
    model.add(MaxPooling2D(pool_size=(2, 2)))
    model.add(Conv2D(20, (3, 3)))
    model.add(Flatten())
    model.add(Dense(units=100, activation='relu'))
    model.add(Dense(units=10, activation='softmax'))
    model.compile(loss=keras.losses.SparseCategoricalCrossentropy(), metrics=['accuracy'],
                  optimizer=keras.optimizers.Adam(lr=0.001))
    return model


def load_test_data():
    (_, _), (x, y) = mnist.load_data()
    l = len(y)
    for i in range(l):
        y[i] = 9 - y[i]
    x = x.reshape(-1, 28, 28, 1)
    return x, y

def train_models():
    (x, y), (_, _) = mnist.load_data()
    x = x.reshape(-1, 28, 28, 1)
    for i in range(participant_number):
        model = new_model()
        model.fit(x, y, batch_size=64, epochs=10)
        model.save("model/"+str(i))

def aggregation(parameters):
    print('aggregation')
    weights = []
    for layer in parameters:
        sum = 0
        l = len(layer)
        for temp in layer:
            sum = sum + temp
        weights.append(sum / l)
    model = new_model()
    l = len(model.get_weights())
    model.set_weights(weights)
    return model

def test(model):
    print('test')
    my_x, my_y = load_test_data()
    loss, acc = model.evaluate(my_x, my_y, batch_size=64)
    if acc > 0.95:
        print("FLAG GET!!!")
    else:
        print("you fail", acc)

def load_parameters():
    print('load parameter')
    parameters = []
    models = []
    for i in range(participant_number):
        models.append(load_model("model/"+str(i)))
    for i in range(8):
        layer = []
        for j in range(participant_number):
            temp = models[j].get_weights()
            layer.append(temp[i])
        parameters.append(layer)
    return parameters

def get_input_parameter_auto(parameters):
    print('get input parameter')
    myid = 0
    for layer in parameters:
        print("next layer:")
        print(layer[0].shape," ",mylayers[myid].shape)
        input_weight = mylayers[myid]
        myid = myid + 1
        layer.append(input_weight)
    return parameters

    
if __name__ == '__main__':
    if not os.path.exists('model'):
        os.mkdir("model")
        train_models()
        
	dict_shape = {0:(3, 3, 1, 10), 1:(10,), 2:(3, 3, 10, 20), 3:(20,), 
                  4:(2420, 100), 5:(100,), 6:(100, 10), 7:(10,)} #每一层的权重结构
	mylayers = []
	for inde in range(8):
		f = open(f'{inde}.txt',mode = 'r')
		file = f.readlines()
		mylayer = []
		for item in file:
			item = item.replace('\n','')
			item = item.split(' ')
			mylayer.append(item)
		mylayer = mylayer[0]
		mylayer = mylayer[0:-1]
		for i in range(len(mylayer)):
			mylayer[i] = float(mylayer[i])
		pas = []
		for num in range(len(dict_shape[inde])):
			pas.append(int(dict_shape[inde][num]))
		if (len(pas)==4):
			mylayers.append(np.array(mylayer).reshape(pas[0],pas[1],pas[2],pas[3]))
		elif(len(pas)==2):
			mylayers.append(np.array(mylayer).reshape(pas[0],pas[1]))
		elif(len(pas)==1):
			mylayers.append(np.array(mylayer))
		print(mylayers[inde].shape)
	f.close()
	parameters = load_parameters()
	parameters = get_input_parameter_auto(parameters)
	model = aggregation(parameters)
	test(model)