基于ember数据集的恶意代码深度学习

2018年4月份，网络安全公司Endgame发布了一款名为EMBER的大型开源数据集。EMBER是一个包含了100多万种良性和恶意PE文件（Windows可执行文件）集合，这是一种常见的恶意软件隐藏格式
笔者刚开始学习相关的知识，根据自己的理解以pytorch为框架复现了论文
Saxe J，Berlin K. Deep neural network based malware detection using two dimensional binary program features[C] // 2015 10th International Conference on Malicious and Unwanted Software(MALWARE). IEEE, 2015: 11-20
中的有关代码，并对ember数据集中“train_features_1.jsonl”的十万余条数据进行训练，希望能对和我一样的小白一些参考，并希望得到专业人士及大佬的一些指正。
首先ember数据集可以在ember数据集中下载
论文中模型的主要框架为

论文一共对二进制的pe文件进行了四组数据的提取，分别是字节熵、直方图、导入函数表、pe元数据信息。
首先是读取json文件

import jsonlines
class ReadData():
    def __init__(self, data, num):
        self.data = data
        self.data_list = []
        self.num = num

    def split_data(self):
        with open(self.data, 'r', encoding='utf-8') as f:
            # 将包含多个json的json文件分隔开
            for item in jsonlines.Reader(f):
                self.data_list.append(item)
            # 取label
        datas = []
        labels = []
        data_list = self.data_list
        for data in data_list[:self.num]:
            if data['label'] >= 0:
              datas.append(data)  
              labels.append(data['label'])
        return datas, labels

利用jsonlines.Reader将数据集中百万余条json数据分隔开存入data_list中，取前num条，因为论文中是做的二分类，而ember数据集中共有三类label：-1，0，1；因此我在这里将label=-1的数据手动去除，将前num条数据中剩下的数据存入datas与labels中
加载数据集

texts, labels = ReadData('train_features_1.jsonl', 200000).split_data()

len(texts),len(labels),len(set(labels))
(116051, 116051, 2)

代表了前20w数据中label为0和1的只有116051条

字节熵

字节熵被用于衡量系统的混乱程度，熵值越大，说明混乱程度越高。我们把PE⽂件当作⼀个字节组成的数组，在这个数组上以2048字节为窗⼝，以1024字节为步长同计算信息熵的方式计算字节熵。
在ember数据集中，每个数据的byteetropy都已经被计算好了，具体数据集是怎样计算的可以参考这位博主的博客：详解EMBER数据集中对PE文件提取ByteEntropyHistogram特征，我在这里只做特征提取，不对其来源做过多的探讨。

import numpy as np
# 第一个类，提取字节熵
class FirstFeature():
    def __init__(self, texts):
        self.texts = texts
    def byteentropy(self):
        entropy = []
        for text in self.texts:
            entropy.append(text['byteentropy'])
        return np.array(entropy)

从每条数据中提取byteentropy
查看shape

entorpy = FirstFeature(texts).byteentropy()

entorpy.shape
(116051, 256)

导入函数表

论文中的意思大概是将PE头的IAT表里面的文件名和函数名hash到0到255范围，再创建一个256维的数组a，每有一个hash后的数num就在这个创建的数组a[num] = a[num] + 1

import hashlib
class SecFeature():
  def __init__(self, texts):
    self.texts = texts
    self.imports_name = self.imports()
  # 取出texts中所有的import名字
  def imports(self):
    imports_name = []
    for text in texts:
      imports_name.append(wash(text['imports']))
    return imports_name
  # 将一个字符串hash到0-256
  def custom_hash(self, input_str):
    hash_object = hashlib.sha256(input_str.encode())  # 使用SHA-256哈希算法
    hash_hex = hash_object.hexdigest()  # 将哈希值转换为十六进制字符串
    hash_int = int(hash_hex, 16)  # 将十六进制字符串转换为整数
    return hash_int
  # 将imports方法中得到的n*len(name)转成n*len(name_to_256)再转成n*256
  def hash_to_256(self):
    name_to_256 = []
    for imp_name in self.imports_name:
      every_num = []
      for single_name in imp_name:
        every_num.append(self.custom_hash(single_name)%256)
      name_to_256.append(every_num)

    all_hash256 = []
    for nums in name_to_256:
      #len(nums)为每个pe文件含有的imports_name数量
      hash256 = [0]*256
      for num in nums:
        hash256[num] = hash256[num] + 1
      all_hash256.append(hash256)
    return np.array(all_hash256)

从数据中提取imports中所有的内容，为了便于分开每个名字，我将其中的无用符号都删去，再custom_hash函数中将字符串hash到整数，之后再进行%256的操作将数值映射到0-255

字节直方图

特征直⽅图本质上PE⽂件也是⼆进制⽂件，可以当作⼀连串字节组成的⽂件。字节直⽅图⼜称为ByteHistogram，它的核⼼思想是，定义⼀个长度为256维的向量，每个向量依次为0x00，0x01⼀直到0xFF，分别代表PE⽂件中0x00，0x01⼀直到0xFF对应的个数。例如经过统计，0x01有两个，0x03和0x05对应的各⼀个，假设直⽅图维度为8，所以对应的直⽅图为：
[0,2,0,1,0,0,1,0]
Python中实现⾃⼰直⽅图⾮常⽅便，numPy提供了⼀个⾮常强⼤的函数：

numpy.bincount(x,weights=None,minlength=None)

实际使⽤时，单纯统计直⽅图⾮常容易过拟合，因为字节直⽅图对于PE⽂件的⼆进制特征过于依赖，PE⽂件增加⼀个⽆意义的0字节都会改变直⽅图；另外PE⽂件中不同字节的数量可能差别很⼤，数量占优势的字节可能会⼤⼤弱化其他字节对结果的影响，所以需要对直⽅图进⾏标准化处理。⼀种常见的处理⽅式是，增加⼀个维度的变量，⽤于统计PE⽂件的字节总数，同时原有直⽅图按照字节总数取平均值。
提取特征
而在ember数据集中已经提取好了字节直方图，所以只需要像字节熵那样直接提取就好

# 第三个类，提取字节直方图
class ThiFeature():
    def __init__(self, texts):
        self.texts = texts
    def get_histogram(self):
        histogram = []
        for text in texts:
            histogram.append(text['histogram'])
        return np.array(histogram)

PE元数据

在论文中提到了PE元数据这个概念，并且提到了时间戳在PE元数据里，而时间戳在ember数据集中的header里，我认为这里提到的PE元数据就是PE文件头信息。

这样的话就只需要像提取导入函数表一样对每个字符串进行hash就好了

class FourFeature():
  def __init__(self, texts):
    self.texts = texts
    self.header_name = self.header()
  # 取出texts中所有的pe头信息
  def header(self):
    header_name = []
    for text in texts:
      header_name.append(wash(text['header']))
    return header_name
  # 将一个字符串hash到0-256
  def custom_hash(self, input_str):
    hash_object = hashlib.sha256(input_str.encode())  # 使用SHA-256哈希算法
    hash_hex = hash_object.hexdigest()  # 将哈希值转换为十六进制字符串
    hash_int = int(hash_hex, 16)  # 将十六进制字符串转换为整数
    return hash_int
  # 将imports方法中得到的n*len(name)转成n*len(name_to_256)再转成n*256
  def hash_to_256(self):
    name_to_256 = []
    for imp_name in self.header_name:
      every_num = []
      for single_name in imp_name:
        every_num.append(self.custom_hash(single_name)%256)
      name_to_256.append(every_num)

    all_hash256 = []
    for nums in name_to_256:
      #len(nums)为每个pe文件含有的pe头字符串数量
      hash256 = [0]*256
      for num in nums:
        hash256[num] = hash256[num] + 1
      all_hash256.append(hash256)
    return np.array(all_hash256)

至此，分别调用四种类中的方法，就得到了四个256维的矩阵，将它们拼接

import torch
entorpy = FirstFeature(texts).byteentropy()
hash256 = SecFeature(texts).hash_to_256()
histogram = ThiFeature(texts).get_histogram()
header_hash = FourFeature(texts).hash_to_256()

entorpy = torch.tensor(entorpy,dtype=torch.float32)
hash256 = torch.tensor(hash256,dtype=torch.float32)
histogram = torch.tensor(histogram,dtype=torch.float32)
header_hash = torch.tensor(header_hash,dtype=torch.float32)
features = torch.cat([entorpy,hash256,histogram,header_hash],dim=1)

features.shape
torch.Size([116051, 1024])

模型

论文中提到了两层隐藏层，一层输出层，还提到了将数据先取以十为底的对数和对每个层根据层的输入和输出的大小进行归一化，在我的试验中，这种方法确实会提高收敛速度和提高精度。
模型主要部分：

    def forward(self, x, label = None):
        x = torch.log10(1+x) # 数据+1防止不能取对数
        
        x = self.linear1(x)
        x = self.drop1(x)
        x = self.batch1(x)
        x = self.prelu1(x)

        x = self.linear2(x)
        x = self.drop2(x)
        x = self.batch2(x)
        x = self.prelu2(x)

        z = self.sigmold(x)
        x = self.clas(x)
        p = self.soft(x)
        self.pre = torch.argmax(p, dim=-1).detach().cpu().numpy().tolist()
        if label is not None:
          loss = self.loss(p, label)
          return loss

归一化

def _weights_init(m):
  if isinstance(m, nn.ReLU or nn.Linear or nn.Sigmoid or nn.Softmax or nn.Tanh or nn.PReLU):
    nn.init.xavier_uniform_(m.weight)
    m.bias.data.zero_()
  elif isinstance(m, nn.BatchNorm1d):
    nn.init.uniform_(m, 0, 0.01)
for param in model.parameters():
  _weights_init(param)

简单的跑十分钟，最后的准确率大概是在96%/93.5%

即是有些过拟合，精度也不是特别很高，但对我来说已经是非常不错的结果了。
最后若有dalao看到这篇文章有任何建议还请不吝赐教

参考文章
1.[当人工智能遇上安全] 4.基于机器学习的恶意代码检测技术详解
2.详解EMBER数据集中对PE文件提取ByteEntropyHistogram特征
3.机器学习⽅法检测恶意⽂件