add injection

utils/comp_two_files.py

@@ -0,0 +1,75 @@
+from bitstring import BitArray
+import os
+
+def compare_bits(file1, file2, max_diff_display=20):
+    """
+    比较两个文件的bit级差异
+    :param file1: 第一个文件路径
+    :param file2: 第二个文件路径
+    :param max_diff_display: 最大差异位置显示数量
+    :return: 差异统计字典
+    """
+    # 读取文件并转换为bit数组
+    bits1 = BitArray(filename=file1)
+    bits2 = BitArray(filename=file2)
+    
+    # 获取bit长度
+    len1, len2 = len(bits1), len(bits2)
+    min_len = min(len1, len2)
+    
+    # 统计结果
+    diff_stats = {
+        'total_bits_file1': len1,
+        'total_bits_file2': len2,
+        'differing_bits': 0,
+        'diff_positions': [],
+        'bit_length_mismatch': len1 != len2
+    }
+    
+    # 逐bit比较
+    for i in range(min_len):
+        if bits1[i] != bits2[i]:
+            diff_stats['differing_bits'] += 1
+            if len(diff_stats['diff_positions']) < max_diff_display:
+                diff_stats['diff_positions'].append(i)
+    
+    # 处理长度不一致的情况
+    if len1 != len2:
+        diff_stats['extra_bits'] = abs(len1 - len2)
+    else:
+        diff_stats['extra_bits'] = 0
+    
+    return diff_stats
+
+def print_diff_report(diff_stats):
+    """打印差异报告"""
+    print(f"Bit长度比较:")
+    print(f"  File1: {diff_stats['total_bits_file1']} bits")
+    print(f"  File2: {diff_stats['total_bits_file2']} bits")
+    
+    if diff_stats['bit_length_mismatch']:
+        print(f"\n! 文件长度不一致，相差 {diff_stats['extra_bits']} bits")
+    
+    print(f"\n差异bit总数: {diff_stats['differing_bits']}")
+    
+    if diff_stats['differing_bits'] > 0:
+        print(f"\n前 {len(diff_stats['diff_positions'])} 个差异位置 (0-based):")
+        for pos in diff_stats['diff_positions']:
+            print(f"  Bit位置 {pos}")
+
+if __name__ == "__main__":
+    # 使用示例
+    file1 = "../malwares/generated_malware"
+    file2 = "../malwares/generated_malware_extracted"
+    
+    # 比较文件
+    diff_stats = compare_bits(file1, file2)
+    
+    # 打印报告
+    print_diff_report(diff_stats)
+    
+    # 高级用法：直接访问差异数据
+    print("\n高级访问:")
+    print(f"总差异bit数: {diff_stats['differing_bits']}")
+    if diff_stats['differing_bits'] > 0:
+        print(f"第一个差异位置: {diff_stats['diff_positions'][0]}")

utils/paraProcess.py

@@ -0,0 +1,80 @@
+"""
+处理参数，如参数翻转、参数替换等等
+"""
+import os
+import torch
+import random
+import struct
+import pandas as pd
+from bitstring import BitArray
+from concurrent.futures import ThreadPoolExecutor
+import numpy as np
+import math
+
+
+def split_file(file_path, chunk_size=8):
+    """
+    分割源文件成为元素BitArray的list
+    :param file_path: 源文件路径
+    :param chunk_size: 分割粒度 Bit
+    :return: 返回一个元素BitArray的list
+    """
+    # 以bit的形式读取文件
+    bit_data = BitArray(filename = file_path)
+    chunks = [bit_data[i:i+chunk_size] for i in range(0, len(bit_data), chunk_size)]
+    return chunks
+
+
+def merge_file(output_file, chunks):
+    """
+    将BitArray的list合并成一个文件
+    :param output_file: 目标文件路径
+    :param chunks: BitArray的list
+    :return: 合并后的文件
+    """
+    merge_data = BitArray()
+    for chunk in chunks:
+        merge_data.append(chunk)
+
+    with open(output_file, 'wb') as file:
+        merge_data.tofile(file)
+
+
+
+def layer_low_n_bit_fLip(initParaPath, flipParaPath, bit_n, *layers):
+    """
+    翻转pth的layers层fa的低n bit
+    :param initParaPath: 原始参数pth
+    :param flipParaPath: 翻转之后的参数pth
+    :param bit_n: 翻转低多少bit
+    :return: void
+    """
+    para = torch.load(initParaPath)
+
+    for layer in layers:  # layers数组中的所有layer
+        if len(para[layer].data.shape) < 1:
+            continue  # 单值除去
+        # print(layer, type(layer))
+        layerTensor = para[layer].data
+        # print(layerTensor.shape)
+        layerTensor_initView = layerTensor.view(torch.int32)
+        # print(format(layerTensor_initView[0][0][0][0], '032b'), layerTensor[0][0][0][0])
+        layerTensor_embedded_int = layerTensor_initView ^ bit_n
+        layerTensor_embedded = layerTensor_embedded_int.view(torch.float32)
+        # print(format(layerTensor_embedded_int[0][0][0][0], '032b'), layerTensor_embedded[0][0][0][0])
+
+        para[layer].data = layerTensor_embedded
+
+    torch.save(para, flipParaPath)
+    return
+
+
+
+
+
+
+
+
+if __name__ == "__main__":
+    
+    print("Test Done")

utils/paraShow.py

@@ -0,0 +1,163 @@
+"""
+展示参数分布、属性
+"""
+import os
+import torch
+import random
+import struct
+import pandas as pd
+from bitstring import BitArray
+from concurrent.futures import ThreadPoolExecutor
+import numpy as np
+import math
+
+from numpy.ma.core import shape
+
+
+def showDif(file1, file2):
+    """
+    对比提取的恶意软件和原始恶意软件的区别,返回出错的bit
+    :return:
+    """
+    malwareStr1 = BitArray(filename=file1).bin
+    malwareStr2 = BitArray(filename=file2).bin
+    diffNum = 0
+    if  malwareStr1 != malwareStr2:
+        print("两个恶意软件大小不同，第一个Bit数为：", malwareStr1, " 第二个Bit数为：", malwareStr1)
+        return
+    for i in range(len(malwareStr1)):
+        if malwareStr1[i] != malwareStr2[i]:  # 打印出所有不同的bit的位置
+            print("pos:", i, "initBit:", malwareStr1[i], "extractedBit:", malwareStr2[i])
+            diffNum += 1
+    # print(malwareStr1)
+    # print(malwareStr2)
+    print("different bit Num between the two files: ", diffNum)
+    return diffNum
+
+
+def get_file_bit_num(file_path):
+    """
+    通过文件路径，获得文件bit数
+    """
+    return os.path.getSize(file_path) * 8
+
+
+def getExpEmbeddSize(initParaPath, layers, interval=1, correct=1):
+    """
+    返回指数部分最大的嵌入容量，单位是字节Byte
+    :param initParaPath:
+    :param layers: list
+    :param interval: 每interval个中嵌入一个
+    :return: list
+    """
+    para = torch.load(initParaPath, map_location=torch.device("cpu"))
+    ret = []
+    for layer in layers:
+        paraTensor = para[layer].data
+        paraTensor_flat = paraTensor.flatten()
+        # print(initParaPath, layers, paraTensor_flat.size())
+        layerSize = len(paraTensor_flat) // (interval * correct * 8)
+        # print(layer, len(paraTensor_flat), layerSize)
+        ret.append(layerSize)
+    return ret
+
+
+def generate_file_with_bits(file_path, num_bits):
+    """
+    根据需要多少bit，随机生成对应大小的恶意软件
+    :param file_path:
+    :param num_bits:
+    :return:
+    """
+    # 计算需要的字节数，每字节有8个bit
+    num_bytes = (num_bits + 7) // 8  # 向上取整，保证比特数足够
+    print("Byte Num:", num_bytes)
+
+    # 创建一个包含随机字节的字节数组
+    byte_array = bytearray(random.getrandbits(8) for _ in range(num_bytes))
+
+    # 如果不需要最后一个字节的全部位，将多余的位清零
+    if num_bits % 8 != 0:
+        last_byte_bits = num_bits % 8
+        # 保留最后字节所需的位数，其它位清零
+        mask = (1 << last_byte_bits) - 1
+        byte_array[-1] &= mask
+
+    # 将字节数组写入文件
+    with open(file_path, 'wb') as f:
+        f.write(byte_array)
+
+    print(f"File '{file_path}' generated with {num_bits} bits.")
+
+
+
+
+
+
+
+if __name__ == "__main__":
+    """测试路径"""
+    path_swin = "../parameters/classification/swin/swin_b.pth"
+    path_yolo = "../parameters/detection/yolov10/yolov10b.pt"
+    path_rt = "../parameters/detection/rt_dert/rt.pth"
+    path_sam = "../parameters/segmentation/samv2/sam.pth"
+    
+    
+    
+    
+    """测试swin"""
+    # swin_keys = torch.load(path_swin).keys()
+    # print(type(torch.load(path_swin)))
+    
+    
+    
+
+    pth_keys = torch.load(path_yolo)
+    print(pth_keys.keys())
+    print(type(pth_keys['model'].model.named_modules()))
+    print(pth_keys['model'].model.named_modules())
+    # for idx, layer in enumerate(pth_keys['model'].model):
+    #     print(f"层 {idx}: {layer}")
+
+
+    # 遍历模型中所有的子模块（包括嵌套层）
+    for idx, (name, module) in enumerate(pth_keys['model'].model.named_modules()):
+        print(f"模块索引 {idx} - 名称 {name}: {module}")
+
+
+
+
+
+
+
+    # print((pth_keys['model'].model[23].cv2[0]))
+    # print((pth_keys['model']))
+    # print(len(pth_keys['model'].model.conv.weight))
+    # print(type(pth_keys['model'].model[1]))
+    # print(type(pth_keys['model'].model[0].conv.weight))
+    # print(type(pth_keys['model'].model[0].conv.weight.data))
+    # print(shape(pth_keys['model'].model[0].conv.weight.data))
+    # print(pth_keys['model'].model[0].conv.weight.data[0][0][0][0].dtype)
+
+
+
+    # path3 = "../parameters/detection/rt_dert/rt.pth"
+    # a = torch.load(path3)
+    # print(a['ema'].keys())
+    # print(a['ema']['module']['backbone.conv1.conv1_1.conv.weight'][0][0][0][0].dtype)
+    #
+    # path4 = "../parameters/segmentation/samv2/sam.pth"
+    # b = torch.load(path4)
+    # print(b.keys())
+    # print(b['image_encoder.neck.0.weight'][0][0][0][0].dtype)
+
+
+
+
+    # print(get_file_bit_num(path1))
+    
+    # print(pth_keys['train_args'])
+    
+    
+    print("Test Done")
+    

utils/swin.py

@@ -0,0 +1,298 @@
+import os
+import torch
+import random
+from bitstring import BitArray
+
+# from utils import generate_file_with_bits, showDif
+
+
+
+def generate_file_with_bits(file_path, num_bits):
+    """
+    根据需要多少bit，随机生成对应大小的恶意软件
+    :param file_path:
+    :param num_bits:
+    :return:
+    """
+    # 计算需要的字节数，每字节有8个bit
+    num_bytes = (num_bits + 7) // 8  # 向上取整，保证比特数足够
+    print("Byte Num:", num_bytes)
+
+    # 创建一个包含随机字节的字节数组
+    byte_array = bytearray(random.getrandbits(8) for _ in range(num_bytes))
+
+    # 如果不需要最后一个字节的全部位，将多余的位清零
+    if num_bits % 8 != 0:
+        last_byte_bits = num_bits % 8
+        # 保留最后字节所需的位数，其它位清零
+        mask = (1 << last_byte_bits) - 1
+        byte_array[-1] &= mask
+
+    # 将字节数组写入文件
+    with open(file_path, 'wb') as f:
+        f.write(byte_array)
+
+    print(f"File '{file_path}' generated with {num_bits} bits.")
+
+
+
+
+class swin:
+    def __init__(self, path):
+        """
+        初始化使用参数的路径进行初始化
+        :param path:
+        """
+        self.path = path
+        return
+
+
+    def get_pth_keys(self):
+        """
+        返回参数的key
+        :param paraPath: 待获得的参数pth
+        :return:
+        """
+        return torch.load(self.path, map_location=torch.device("cpu")).keys()
+
+
+    def get_pth_keys_float32(self):
+        """
+        返回参数的key
+        :param paraPath: 待获得的参数pth
+        :return:
+        """
+        para = torch.load(self.path, map_location=torch.device("cpu"))
+        temp = para.keys()
+        layers = []
+        for i in temp:
+            if para[i].data.dtype == torch.float32:
+                layers.append(i)
+        return layers
+
+
+    def get_file_bit_num(self):
+        """
+        通过文件路径，获得文件bit数
+        :return: bit size
+        """
+        return os.path.getSize(self.path) * 8
+
+
+    def layer_low_n_bit_fLip(self, flip_path, bit_n, *layers):
+        """
+        翻转pth的layers层的低n bit
+        :param flip_path: 翻转之后的参数pth
+        :param bit_n: 翻转低多少bit
+        :return: void
+        """
+        para = torch.load(self.path)
+        mask= (1<<bit_n)-1
+        for layer in layers:
+            if len(para[layer].data.shape) < 1:
+                continue
+            layer_tensor = para[layer].data
+            # 确保 Tensor 的数据类型为浮点型
+            if layer_tensor.dtype == torch.float32:
+                # 使用整数视图来操作比特，但应谨慎操作
+                layer_tensor_view = layer_tensor.view(torch.int32)
+                layer_tensor_view ^= mask  # 对低位进行翻转操作
+                para[layer].data = layer_tensor_view.view(torch.float32)
+        torch.save(para, flip_path)
+
+
+    def all_layers_low_n_bit_fLip(self, flip_path, bit_n):
+        """
+        翻转所有层的低n bit，需要满足其中的数据类型是fp32
+        :param flip_path: 翻转之后的参数pth
+        :param bit_n: 翻转低多少bit
+        :return:
+        """
+        self.layer_low_n_bit_fLip(flip_path, bit_n, *self.get_pth_keys_float32())
+
+
+    def get_layers_low_n_bit_size(self, layers, bit_n):
+        """
+        usage:
+        size, size_list = agent.get_layers_low_n_bit_size(agent.get_pth_keys(), 16)
+        返回指数部分最大的嵌入容量，单位是字节bit
+        :param layers: list
+        :param interval: 每interval个中嵌入一个
+        :return: 总大小，每一层的大小
+        """
+        para = torch.load(self.path, map_location=torch.device("cpu"))
+        size_with_layers_list = []
+        total_size = 0
+        for layer in layers:
+
+            if para[layer].data.dtype != torch.float32:
+                continue  # 只查看是float32数据类型的数值进行嵌入
+
+            paraTensor = para[layer].data
+            paraTensor_flat = paraTensor.flatten()
+            layerSize = len(paraTensor_flat) * bit_n
+            size_with_layers_list.append(layerSize)
+            total_size += layerSize
+        return total_size, size_with_layers_list
+
+
+
+    def all_layers_low_n_bit_inject(self, inject_path, bit_n, malware, malware_len):
+        """
+        随机生成一个软件、在所有层的低nbit进行嵌入
+        :param inject_path: 最后的嵌入pth文件
+        :param bit_n: 低nbit
+        :param malware: 需要嵌入的恶意软件
+        :param malware_len: 需要嵌入的恶意软件的长度，单位为bit
+        :return:
+        """
+        paras = torch.load(self.path, map_location=torch.device("cpu"))
+        malware_str = BitArray(filename=malware).bin
+        mal_index = 0  # 需要访问的恶意软件的bit区间起点
+        _, size_list = self.get_layers_low_n_bit_size(self.get_pth_keys_float32(), bit_n)
+        for layer, size in zip(self.get_pth_keys_float32(), size_list):
+            if paras[layer].data.dtype != torch.float32:
+                continue # 只考虑32位浮点数
+            print(layer, size)
+            para_tensor_flat = paras[layer].flatten()
+            # index 再恶意软件中和tensor中并不是同一个！
+            print(mal_index, mal_index + size)
+            print(para_tensor_flat.size(), size//bit_n)
+            para_index = 0  # 写入的参数的位置
+            for inject_pos in range(mal_index, min(mal_index + size, malware_len), bit_n):
+                current_write_content = malware_str[inject_pos: inject_pos + bit_n]  # 将n bit数据提取出来
+                para_tensor_flat_str = BitArray(int=para_tensor_flat[para_index].view(torch.int32), length=32).bin
+                new_para_tensor_flat_str = para_tensor_flat_str[:32 - bit_n] + current_write_content
+
+                if int(new_para_tensor_flat_str, 2) >= 2 ** 31:
+                    newParaInt = torch.tensor(int(new_para_tensor_flat_str, 2) - 2 ** 32, dtype=torch.int32)
+                    para_tensor_flat[para_index] = newParaInt.view(torch.float32)
+                else:
+                    newParaInt = torch.tensor(int(new_para_tensor_flat_str, 2), dtype=torch.int32)
+                    para_tensor_flat[para_index] = newParaInt.view(torch.float32)
+                para_index += 1  # 写入的位置往后推1bit
+
+            if mal_index + size >= malware_len:
+                break
+            else:
+                mal_index = mal_index + size
+            paras[layer] = para_tensor_flat.reshape(paras[layer].data.shape)
+        torch.save(paras, inject_path) 
+        return
+
+
+    def all_layers_low_n_bit_extract(self, inject_path, bit_n, extract_malware, malware_len):
+        """
+        :param malware_len: 需要嵌入的恶意软件的长度，单位为bit
+        :param inject_path: 嵌入后的模型参数路径
+        :param bit_n: 嵌入参数的后nbit
+        :param extract_malware: 提取出来的恶意软件路径
+        :param malware_len: 需要嵌入的恶意软件的长度，单位为bit
+        :return:
+        """
+        paras = torch.load(inject_path, map_location="cpu"); bits, idx = BitArray(), 0;
+        _, size_list = self.get_layers_low_n_bit_size(self.get_pth_keys_float32(), bit_n)
+        for layer, _ in zip(self.get_pth_keys_float32(), size_list):
+            p = paras[layer].data
+            if p.dtype != torch.float32: continue
+            for x in p.flatten()[:min(len(p.flatten()), (malware_len - idx + bit_n - 1) // bit_n)]:
+                bits.append(f'0b{BitArray(int=int(x.view(torch.int32)), length=32).bin[-bit_n:]}');
+                idx += bit_n
+                if idx >= malware_len: break
+            if idx >= malware_len: break           
+        with open(extract_malware, 'wb') as f:
+            bits_clip = bits[:(malware_len-(malware_len%bit_n))] + bits[-(malware_len%bit_n):]
+            bits_clip[:malware_len].tofile(f)
+        return
+
+
+    # def all_layers_low_n_bit_extract(self, inject_path, bit_n, extract_malware, malware_len):
+    #     paras = torch.load(inject_path, map_location="cpu");
+    #     pl = torch.load(self.path, map_location="cpu");
+    #     layers = [k for k, v in pl.items() if v.dtype == torch.float32];
+    #     bits, idx = BitArray(), 0
+    #     for l in layers:
+    #         f = paras[l].data.flatten();
+    #         r = min(len(f), (malware_len - idx + bit_n - 1) // bit_n)
+    #         for x in f[:r]:
+    #             bits.append(f'0b{BitArray(int=int(x.view(torch.int32)), length=32).bin[-bit_n:]}');
+    #             idx += bit_n
+    #             if idx >= malware_len: break
+    #         if idx >= malware_len: break
+    #     with open(extract_malware, 'wb') as f:
+    #         bits[:malware_len].tofile(f); return
+
+
+if __name__ == "__main__":
+    
+    
+    
+    path = "../parameters/classification/swin_face/swin_face.pth"
+    # flip_path = "../parametersProcess/swin/swin_flip_16.pth"
+    inject_path = "../parametersProcess/swin_face/swin_evilfiles_16.pth"
+    malware = "../malwares/generated_malware"
+    extract_malware = "../malwares/generated_malware_extracted"
+    agent = swin(path)
+
+
+    # print("layers name: ", agent.get_pth_keys_float32())
+    # print("type: ", type(agent.get_pth_keys_float32()))
+    # print("layers num: ", len(agent.get_pth_keys_float32()))
+
+    size, size_list = agent.get_layers_low_n_bit_size(agent.get_pth_keys_float32(), 16)
+    # print("all layers injection size with low-16 bits: ", size / 8000000, " MB")
+    # print(size_list)
+    # print(len(size_list))
+    '''随机生成一个恶意软件，全部嵌入模型的层（简化流程）'''
+    generate_file_with_bits(malware, size)
+
+
+    print("malware bit size: ",os.path.getsize(malware) * 8)
+
+
+    '''嵌入'''
+    agent.all_layers_low_n_bit_inject(inject_path, 20, malware, os.path.getsize(malware) * 8)
+
+    '''提取'''
+
+
+    # def all_layers_low_n_bit_extract(ip, bn, em, ml):
+    #     p = torch.load(ip, map_location="cpu");
+    #     b, i = BitArray(), 0;
+    #     lrs = [k for k, v in p.items() if v.dtype == torch.float32]
+    #     for l in lrs:
+    #         for x in p[l].data.flatten()[:min(len(p[l].data.flatten()), (ml - i + bn - 1) // bn)]:
+    #             b.append(f'0b{BitArray(int=int(x.view(torch.int32)), length=32).bin[-bn:]}');
+    #             i += bn
+    #             if i >= ml: break
+    #         if i >= ml: break
+    #     with open(em, 'wb') as f:
+    #         b[:ml].tofile(f);return
+
+
+    agent.all_layers_low_n_bit_extract(inject_path, 20, extract_malware, os.path.getsize(malware) * 8)
+    # agent.all_layers_low_n_bit_extract(inject_path, 16, extract_malware, 36272)
+
+
+    # def all_layers_low_n_bit_extract(ip, bn, em, ml):
+    #     p = torch.load(ip, map_location="cpu");
+    #     b, i = BitArray(), 0;
+    #     lrs = [k for k, v in p.items() if v.dtype == torch.float32]
+    #     for l in lrs:
+    #         for x in p[l].data.flatten()[:min(len(p[l].data.flatten()), (ml - i + bn - 1) // bn)]:
+    #             b.append(f'0b{BitArray(int=int(x.view(torch.int32)), length=32).bin[-bn:]}');
+    #             i += bn
+    #             if i >= ml: break
+    #         if i >= ml: break
+    #     with open(em, 'wb') as f:
+    #         b[:ml].tofile(f);return
+    # all_layers_low_n_bit_extract("~/data/ATATK/parametersProcess/swin/swin_inject_16.pth", 16, "~/data/ATATK/malwares/Zherkov_extract.EXE", 36272)
+
+
+    
+
+
+
+
+
+    # agent.all_layers_low_n_bit_fLip(flip_path, 20)