import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.transforms as transforms
import numpy as np
from PIL import Image
import torch.nn.functional as F

class tanji_model(nn.Module):
    def __init__(self):
        super(tanji_model, self).__init__()
        self.classifier = nn.Sequential(
            nn.Linear(31 * 31, 512),
            nn.ReLU(),
            nn.Linear(512, 128),
            nn.ReLU(),
            nn.Linear(128, 1),
        )

    def forward(self, x):
        x = torch.flatten(x, 1)
        y = self.classifier(x)
        return torch.sigmoid(y)

# 加载与训练模型文件model.pth
tanji = tanji_model().cpu()
tanji.load_state_dict(torch.load('model.pth'))
tanji.eval()
print(tanji)

# 创建一个跟flag一样的空白图片
flag_size = (31, 31)
img = torch.zeros((1, 1, * flag_size), requires_grad = True)
# 设置优化器
optimizer = optim.Adam([img], lr=0.001)

for i in range(10000):
    # 把梯度置零
    optimizer.zero_grad()
    output = tanji(img)
    # 计算损失 _ 最小化负对数似然损失
    loss = -torch.log(F.sigmoid(output))
    # 反向传播
    loss.backward()
    # 更新参数
    optimizer.step()
    with torch.no_grad():
        img.clamp_(0, 1)
    if i % 100 == 0:
        print(f'Iteration {i}, Loss: {loss.item()}')

# 保存图片
optimized_image = transforms.ToPILImage()(img.squeeze(0))
# optimized_image = transforms.ToPILImage()(img[0].cpu()).save('flag.png')
optimized_image.show()
optimized_image.save('flag.png')

import torch  
import torch.nn as nn  
import numpy as np  
# 解决报错：OMP: Error #15: Initializing libiomp5md.dll, but found libiomp5md.dll already initialized.  
import os  
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"  
  
# 读取文件并解析密文数据  
ch = []  
with open('ciphertext.txt', 'r') as file:  
    for line in file:  
        try:  
            line_numbers = line.strip().split()  
            ch.append([float(num) for num in line_numbers])  
        except ValueError:  
            print(f"Ignored non-numeric line: {line}")  
print(len(ch))  
n = len(ch)-1  
  
# 定义加密网络  
class Net(nn.Module):  
    def __init__(self):  
        super(Net, self).__init__()  
        self.linear = nn.Linear(n, n*n)  
        self.conv = nn.Conv2d(1, 1, (2, 2), stride=1, padding=1)  
  
    def forward(self, x):  
        x = self.linear(x)  
        x = x.view(1, 1, n, n)  
        x = self.conv(x)  
        return x  
  
# 加载训练好的模型参数  
mynet=Net()  
mynet.load_state_dict(torch.load('model.pth'))  
conv_weight = mynet.conv.weight.data.tolist()[0][0]  
print(conv_weight)  
conv_bias = mynet.conv.bias.data.tolist()  
print(conv_bias)  
linear_weight = mynet.linear.weight.data.tolist()  
print(linear_weight[0])  
linear_bias = mynet.linear.bias.data.tolist()  
print(len(linear_bias))  
  
# 解密  
'''  
初始化 flag 数组，用于存储解密后的数据  
遍历 ch 中的每个元素，通过逆向操作计算出原始数据。具体操作包括减去卷积层的偏置和加权值，然后除以卷积核的权重。  
  
逆卷积操作  
对于卷积层，卷积核的大小为 2x2，并且在正向操作中，卷积公式为：  
output[i][j]=conv_weight[0][0]×input[i][j]+conv_weight[0][1]×input[i][j+1]+conv_weight[1][0]×input[i+1][j]+conv_weight[1][1]×input[i+1][j+1]+conv_bias  
  
将 flag 转换为适当的形状，并减去全连接层的偏置。  
使用全连接层权重的伪逆矩阵对 flag 进行矩阵运算，得到解密后的数据。  
将解密后的数据转换为字符并打印出来，形成解密后的文本。  
'''  
flag = [[0]*n for _ in range(n)]  
for i in range(n):  
    for j in range(n):  
        x = 0  
        y = 0  
        z = 0  
        if i>0:  
            y = flag[i-1][j]  
            if j>0:  
                x = flag[i-1][j-1]  
        if j>0:  
            z = flag[i][j-1]  
        # 减去卷积层的偏置 conv_bias[0] 和邻近元素的加权和，再除以卷积核 conv_weight[1][1]。  
        flag[i][j]=ch[i][j]-conv_bias[0]-(conv_weight[0][0]*x+conv_weight[0][1]*y+conv_weight[1][0]*z)  
        flag[i][j]/=conv_weight[1][1]  
  
# 转换和偏置调整  
'''  
将 flag 转换为一维数组。  
减去全连接层的偏置 linear_bias。  
将 flag 重新调整为 (n*n, 1) 的形状。  
'''  
flag = np.array(flag).reshape(n*n)  
flag = flag-linear_bias  
flag = flag.reshape(n*n,1)  
print(flag.shape)  
print(np.array(linear_weight).shape)  
print(np.linalg.pinv(linear_weight).shape)  
# 使用全连接层权重 linear_weight 的伪逆矩阵对 flag 进行矩阵运算，得到解密后的数据。  
flag = np.linalg.pinv(linear_weight).dot(flag)  
for i in flag:  
    print(chr(round(i[0])),end='')

import binascii  

data = b''  
with open('encrypted.pth', 'rb') as f:  
    data = f.read()  
    #每位异0xde后保存  
    for i in range(256):  
        print(binascii.hexlify(data[:2]).decode('utf-8'))  
        if data[0]^i==0x50 and data[1]^i==0x4B:  
            print('key:', hex(i))  
            data = bytes([j ^ i for j in data])  
            # 判断文件头十六进制是否为50 4B  
            with open('decrypted.pth', 'wb') as f:  
                f.write(data)  
                break