基于机器学习中图像检测的恶意软件分类技术——尝试与实践
(注：本博客中涉及的实验由于小组资源和实力有限，可能缺乏严谨性)
上海交大荣昶杯比赛真的很棒！学到了很多东西！
那么在此次比赛中我们小组负责的项目是基于cnn图像检测的恶意软件识别技术，故写一篇博客记录一下
本次项目github仓库：https://github.com/BSDFZ-programming-team/MalPro
- 实验3.1-提取.asm图像特征和Opcode 3-gram特征训练决策树模型
- 实验3.2-提取.asm图像特征和Opcode 3-gram特征训练随机森林模型
- 实验3.3-提取.asm图像特征和Opcode 3-gram特征训练cnn模型
0x00 写在前面
我们的对象是IDA反汇编工具生成的.asm文件，思路也很明确，将.asm文件转化为灰度图，再通过特征提取识别恶意代码类型
那这就不得不提到微软前几年搞的比赛了：https://www.kaggle.com/competitions/malware-classification
数据集来源：https://www.kaggle.com/competitions/malware-classification/data?select=train.7z
数据集内容：九类不同类型的恶意软件，编号为[1-9]
那么，故事要从恶意软件图像开始说起。。。。
0x01 前置知识
恶意软件图像
这个概念首先由Nataraj和Karthikeyan的论文Malware Images: Visualization and Automatic Classification中提出，
A given malware binary is read as a vector of 8 bit unsigned
integers and then organized into a 2D array. This can be visualized
as a gray scale image in the range [0,255] (0: black, 255: white).
The width of the image is fixed and the height is allowed to vary
depending on the file size (Fig. 1). Tab. 1 gives some recommended
image widths for different file sizes based on empirical
observations.
论文中提取.asm文件的GIST特征进行恶意软件图像生成，使得同一种类的病毒图片具有相似性
那么我们可以写出将.asm文件转化为恶意软件图像的代码，如下：
# https://bindog.github.io/blog/2015/08/20/microsoft-malware-classification
import numpy
from PIL import Image
import binascii
def getMatrixfrom_bin(filename,width):
with open(filename, 'rb') as f:
content = f.read()
hexst = binascii.hexlify(content) #将二进制文件转换为十六进制字符串
fh = numpy.array([int(hexst[i:i+2],16) for i in range(0, len(hexst), 2)]) #按字节分割
rn = len(fh)/width
fh = numpy.reshape(fh[:rn*width],(-1,width)) #根据设定的宽度生成矩阵
fh = numpy.uint8(fh)
return fh
filename = "your_bin_filename"
im = Image.fromarray(getMatrixfrom_bin(filename,512)) #转换为图像
im.save("your_img_filename.png")
论文中指出：
various malware families have distinct visual characteristics.
由此可知，不同的病毒所生成的恶意软件图像具有不同的特征，相同种类的病毒所生成的恶意软件图像具有相似的特征
MalConv概述
在Edward Raff, Jon Barker的Malware Detection by Eating a Whole EXE中指出了一种全新的检测方法，即，通过读取原始文件的每一个字节作为输入特征，使得计算量和内存用量能够根据序列长度而高效地扩展，在检查整个文件的时候能同时考虑到本地和全局上下文，以及在分析标记为恶意软件的时候能够提供更好的解释能力
其神经网络基于rnn，原理如上图
随机森林算法（Random Forests）
Random Forests算法最初由Leo Breiman和Adele Cutler在2001年提出(https://www.stat.berkeley.edu/~breiman/RandomForests/)
随机森林是一个非常强大的机器学习方法，顾名思义，它是用随机的方式建立一个森林，森林里面有很多的决策树组成，随机森林的每一棵决策树之间是没有关联的
这里的估计器就是一个个的决策树，由他们组成“决策树森林”
随机森林算法是以决策树为估计器的Bagging算法
通过每个决策树的结果，“集思广益”。单个决策树对训练集的噪声非常敏感，但通过Bagging算法降低了训练出的多颗决策树之间关联性，有效缓解了问题
Opcode操作码 & N-gram
Opcode操作码：机器语言指令的一部分，操作码的操作可以包括算术、数据操作、逻辑操作和程序控制。
N-gram特征是给定恶意样本中每相邻n个字节的序列模型，可以理解为一阶马尔科夫链，即字节序列的出现顺序是离散事件的随机过程。
基于ByteCode将n-gram应用于恶意代码识别的想法最早由Tony等人在2004年的论文N-gram-based Detection of New Malicious Code 中提出，随后2008年Moskovitch等人的论文Unknown Malcode Detection Using OPCODE Representation 中提出利用OpCode代替ByteCode的方法
由此，我们可以得出一些基本的点子：
1.直接使用cnn对恶意软件图像进行经典的三通道多元图像识别
2.使用MalConv直接对原始二进制文件识别
0x02 初步实验及结论
实验2.1-直接使用cnn对恶意软件图像进行经典的三通道多元图像识别
我们采用最经典的卷积神经网络对恶意软件图像进行识别
训练代码：
train.py
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
import torchvision.transforms as transforms
from torchvision.datasets import ImageFolder
import torch.multiprocessing as mp
# Define the CNN model
class TrafficSignClassifier(nn.Module):
def __init__(self, num_classes):
super(TrafficSignClassifier, self).__init__()
self.conv1 = nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1)
self.relu1 = nn.ReLU()
self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)
self.relu2 = nn.ReLU()
self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
self.fc1 = nn.Linear(64 * 32 * 32, 128)
self.relu3 = nn.ReLU()
self.fc2 = nn.Linear(128, num_classes)
def forward(self, x):
x = self.conv1(x)
x = self.relu1(x)
x = self.pool1(x)
x = self.conv2(x)
x = self.relu2(x)
x = self.pool2(x)
x = x.view(x.size(0), -1)
x = self.fc1(x)
x = self.relu3(x)
x = self.fc2(x)
return x
if __name__ == '__main__':
mp.freeze_support()
# Set up the training and validation data
transform = transforms.Compose([
transforms.Resize((128, 128)),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
train_dataset = ImageFolder(r"./train", transform=transform)
train_loader = DataLoader(train_dataset, batch_size=2, shuffle=True, num_workers=4,drop_last=True)
# Set up the model, loss function, and optimizer
model = TrafficSignClassifier(num_classes=len(train_dataset.classes))
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)
# Train the model
num_epochs = 20
for epoch in range(num_epochs):
train_loss = 0.0
train_correct = 0
train_total = 0
model.train()
for inputs, labels in train_loader:
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
train_loss += loss.item()
_, predicted = torch.max(outputs.data, 1)
train_total += labels.size(0)
train_correct += (predicted == labels).sum().item()
print("训练了"+str(epoch+1)+"次")
torch.save(model.state_dict(), './model.pt')
对最终结果进行测试
EXAM.py
import torch
from PIL import Image
import os
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
import torchvision.transforms as transforms
from torchvision.datasets import ImageFolder
import torch.multiprocessing as mp
from torchvision.transforms import ToTensor, Normalize, Compose, Resize
from pathlib import Path
class_names = [str(i) for i in range(1, 10)]
# 定义图像预处理步骤
transform = Compose([
Resize((64, 64)), # 调整图像大小以匹配模型输入尺寸
ToTensor(), # 将图像转换为Tensor
Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), # 标准化
])
def load_image(image_path):
# 打开图像文件
image = Image.open(image_path).convert('RGB') # 确保图像是RGB格式
# 应用预处理步骤
image = transform(image)
# 添加一个批次维度，因为模型期望输入是批次格式
image = image.unsqueeze(0)
return image
class TrafficSignClassifier(nn.Module):
def __init__(self, num_classes):
super(TrafficSignClassifier, self).__init__()
self.conv1 = nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1)
self.relu1 = nn.ReLU()
self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)
self.relu2 = nn.ReLU()
self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
self.fc1 = nn.Linear(64 * 32 * 32, 128)
self.relu3 = nn.ReLU()
self.fc2 = nn.Linear(128, num_classes)
def forward(self, x):
x = self.conv1(x)
x = self.relu1(x)
x = self.pool1(x)
x = self.conv2(x)
x = self.relu2(x)
x = self.pool2(x)
x = x.view(x.size(0), -1)
x = self.fc1(x)
x = self.relu3(x)
x = self.fc2(x)
return x
model_path = r'./model.pt'
model = TrafficSignClassifier(num_classes=9).to('cpu') # 替换...为实际的类别数
model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))
# 设置模型为评估模式
model.eval()
# 要预测的图片目录
images_dir = r"./test"
correct = 0
total = 0
# 预测并重命名图片
for img_path in Path(images_dir).glob('*'):
# 加载图片
input_image = load_image(img_path)
# 进行模型推理
with torch.no_grad(): # 在推理过程中不需要计算梯度
outputs = model(input_image)
_, predicted_class = torch.max(outputs.data, 1)
# 获取预测类别名称
predicted_label = class_names[predicted_class.item()]
#predicted_label = str(predicted_class.item())
ans = img_path.name.split('_')[0]
if ans == predicted_label:
correct += 1
else:
pass
total += 1
print('ans: '+ans +'predicted: '+ predicted_label)
# new_filename = f"{predicted_label}_{img_path.stem}{img_path.suffix}" #_{img_path.stem}{img_path.suffix}
# new_path = os.path.join(images_dir, new_filename+'.png')
# os.rename(img_path, new_path)
# print(f"Renamed {img_path.name} to {new_filename}")
print(f'{correct}/{total}')
实验过程中，经过多次调参，最终调成上述代码中的参数
实验结果：
最终在epoch=10时，我们得出正确率为103/180，约为53.9%
相较于完全随机猜测的1/9有所提升
实验结论:
1.模型并不能很好地达成识别目的，然而又并不是完全随机猜测；
2.注意到，cnn模型对于2，6，7，8，9四类软件识别效果较好，正确率能达到82%；
3.相反，对于1，3，5识别效果较差，正确率仅为8%左右，出现大量将1，3识别成2的错误情况
实验2.2-使用MalConv直接对原始二进制文件识别
接下来，我们小组尝试直接采用对MalConv对.asm文件进行识别
github项目地址：https://github.com/CaLlMeErIC/MalConv
我们小组修改了项目，将原本的二元检测改为九种不同的恶意软件分类
将数据集中的trainlabels.csv文件转化为json文件
convert.py
import csv
import json
csv_file = 'trainLabels.csv'
json_file = 'label_dict.json'
data = {}
with open(csv_file, mode='r') as file:
csv_reader = csv.reader(file)
for row in csv_reader:
key = row[0]
value = row[1]
data[key] = value
with open(json_file, mode='w') as file:
json.dump(data, file, indent=4)
print(f'JSON file has been created: {json_file}')
接着修改标签相关的内容
读取标签时修改标签形式为1-9即可
小组在实验过程中遇到了算力不足的问题（MALCONV比较吃算力）通过降低参数值解决
实验结果：
单核CPU运行约600s两算法对比（MalConv不保留小数）：
训练过程中损失率始终高于60%
最终正确率约为51%，且各个种类的恶意图像分布均匀
实验结论：
1.Malconv对算力需求较高，降低参数将导致准确率大打折扣
2.Malconv不存在类似cnn中对于图像特征不明显的病毒种类误报率高的问题，对于每一种病毒准确率和错解都比较均匀
0x03 改进&优化
小组决定对第一个实验项目“直接使用cnn对恶意软件图像进行经典的三通道多元图像识别”进行改进
在malware-classification比赛中，我们发现，冠军队伍仅仅提取了.asm恶意软件图像的前800个像素作为特征点进行描述
这就意味着，他们放弃了从.asm文件中提取出来的GIST特征，而去采用前800个像素点的内容作为特征进行提取
例：一个.asm文件的前800个像素
显然，前800个像素对于并不是PE文件里的内容，而是IDA提供的信息。或者说，实际上冠军队伍的方法压根与恶意软件图像没有关系，实际上是用到了IDA产生的信息（虽然这些信息看起来没什么用）。。。
然而除了这种技俩之外，他们还是用了更为先进的算法，即提取opcode 3-gram特征
实验3.1-提取.asm图像特征和Opcode 3-gram特征训练决策树模型
小组通过修改比赛冠军代码项目脚本进行实验
实验结果
将该项目修改为决策树算法后，得到如下准确率
最终准确率约为95.56%
实验结论
提取opcode 3gram和.asm图像特征（前1500px）能大幅度提升训练精度
实验3.2-提取.asm图像特征和Opcode 3-gram特征训练随机森林模型
小组受到实验启发，决定将决策树算法优化为随机森林算法
实验结果：
最终，小组得出了接近99%的正确率
同时，决策树与随机森林单次交叉验证效果对比
实验结论
由此可得，在小样本，小算力的前提下，针对.asm文件的识别及分类，转化为灰度图并不是最佳解法
随机森林在训练效率上优于决策树算法
然而在针对PE文件的识别中，由于失去了asm image features，情况可能有所不同
实验3.3-提取.asm图像特征和Opcode 3-gram特征训练cnn模型
随后，小组决定总结前两个实验的数据，提取出.asm图像特征和Opcode 3-gram特征对实验2.1中的cnn模型进行训练
小组在进行训练的过程中对cnn的数据集和参数进行了一些修改，具体代码参见github仓库（测试处尚未完善，可以先去看一些别人的CNN模型XD）
实验结果
我们得出了令人惊叹的实验结果（我们仍然对此次实验结果存疑（数据太美了 XD））
但是，cnn在速度，效率方面都远远弱于随机森林算法，这也是我们在最终项目中不打算采用cnn的原因
实验结论
由上图，使用cnn算法训练提取出.asm图像特征和Opcode 3-gram特征也许在小算力，小时间，小样本量的前提下能取得比随机森林算法更好的效果（？）
虽然小组尽可能地保证实验过程严谨，但是由于算力问题，以及小组本身实力实在有限，故此结论仅供参考
0x04 实验结论
从上述实验中我们可以得出，提取出.asm图像特征+opcode-3gram特征进行cnn模型的训练能得出更好的结果，但是效率太慢
而随机森林模型不仅具有比较好的效率，在正确率上也相对高
而.asm图像特征+opcode-3gram特征远远强于直接对恶意软件图像进行训练和识别
0x05 参考文献
Malware Detection by Eating a Whole EXE
Malware Images: Visualization and Automatic Classification
利用机器学习进行恶意代码分类
RandomForests
N-gram-based Detection of New Malicious Code
Unknown Malcode Detection Using OPCODE Representation
ToyMalwareClassification(github repo)
MalConv(github repo)
浙公网安备 33010602011771号