基于机器学习的恶意软件检测系统构建与实践

基于机器学习的恶意软件检测系统构建与实践

本文详细介绍了如何使用人工神经网络（ANN）构建一个恶意软件检测系统。从传统恶意软件检测方式的局限性入手，介绍了Windows PE格式的特点，并基于此格式设计了特征工程。文章还详细描述了数据集的划分、模型训练过程以及评估方法。最后，展示了如何将训练好的模型部署为API服务，并通过实际测试验证了其效果。

传统恶意检测方式

传统的恶意软件检测引擎依赖于使用签名，简单理解就是给已知恶意程序生成一个md5数据，使用这个md5来识别是不是恶意程序；很明显，此类检测手段很容易绕过，比如在程序中填充一些数据，生成的md5就会变化，那么就检测不到。

Windows PE格式

Windows PE 格式 官方介绍：https://docs.microsoft.com/en-us/windows/desktop/debug/pe-format

1. PE 具有多个标头，描述其属性和各种寻址细节，例如 PE 将在内存中加载的基地址以及入口点的位置。
2. PE 有多个部分，每个部分包含数据（常量、全局变量等）、代码（在这种情况下该部分被标记为可执行）或有时两者兼有。
3. PE 包含导入了哪些 API 以及从哪些系统库导入的声明

例如，Firefox PE 部分如下所示

而在某些情况下，如果 PE 已经用UPX 之类的打包程序处理过，它的部分可能会看起来有点不同，因为主代码和数据部分被压缩了，并且添加了一个在运行时解压缩的代码存根：

数据集

约 20万 个Windows PE样本，分为恶意样本（VirusTotal 上有 10+ 检测结果）和 安全样本（已知且 VirusTotal 上有 0 检测结果）

由于在同一个数据集上训练和测试模型没有多大意义（因为它可以在训练集上表现得非常好，但无法在新的样本上推广），因此该数据集将自动分为 3 个子集：

• 训练集：包含70%的样本，用于训练。
• 验证集：包含15% 的样本，用于在每个训练阶段对模型进行基准测试。
• 测试集：包含 15% 的样本，用于训练后对模型进行基准测试。

理想情况下，数据集应定期使用较新的样本更新，并重新训练模型，以便即使在出现新的独特样本时也能保持其准确性。

数据预处理特征工程

基于Windows PE格式，将这些本质上非常异构的值（它们是各种间隔的数字和可变长度的字符串）编码为标量数字向量，每个标量数字都在区间 [0.0,1.0] 内标准化，并且长度恒定，变成机器学习模型能够理解的数据

首先需要使用开源的工具ergo（https://github.com/evilsocket/ergo） 来生成一个检测项目

ergo create ergo-pe-sec

根据开源工具的用法，在encoder.py文件中实现特征提取算法，

各特征数据处理

基础属性

根据PE文件格式提取11个特征向量用于判断、检测，为true则即为1.0，为false则写0.0

点击图片可查看完整电子表格

处理代码

def encode_properties(pe):
    global properties
    props = np.array([0.0] * len(properties))
    for idx, prop in enumerate(properties):
        props[idx] = 1.0 if getattr(pe, prop) else 0.0
    return props

PE入口点函数

PE 入口点函数的前 64 个字节，每个元素都[0.0,1.0]通过除以来标准化255这将有助于模型检测那些具有非常独特的入口点的可执行文件，这些入口点在同一家族的不同样本之间仅略有不同（可以将其视为一个非常基本的签名）

ep_bytes=[0]*64
try:
    ep_offset = pe.entrypoint - pe.optional_header.imagebase
    ep_bytes = [int(b) for b in raw[ep_offset:ep_offset+64]]
except Exception as e:
    log.warning("can't get entrypoint bytes from %s: %s", filepath, e)
# ...
# ...
def encode_entrypoint(ep):
    while len(ep) < 64: # pad
        ep += [0.0]
    return np.array(ep) / 255.0 # normalize

重复字节数据特征

二进制文件中 ASCII 表每个字节重复数据 - 该数据点将对有关文件原始内容的基本统计信息进行编码

def encode_histogram(raw):
    histo = np.bincount(np.frombuffer(raw, dtype=np.uint8), minlength=256)
    histo = histo / histo.sum() # normalize
    return histo

PE导入表

手动选择了数据集中最常见的 150 个库（https://github.com/evilsocket/ergo-pe-av/blob/master/encoder.py#L22），并且对于 PE 使用的每个 API，将相关库的列增加一，创建另一个包含 150 个值的直方图，然后根据导入的 API 总量进行归一化

def encode_libraries(pe):
    global libraries
    imports = {dll.name.lower():[api.name if not api.is_ordinal else api.iat_address
        for api in dll.entries] for dll in pe.imports}
    libs = np.array([0.0] * len(libraries))
    for idx, lib in enumerate(libraries):
        calls = 0
        dll= "%s.dll" % lib
        if lib in imports:
            calls = len(imports[lib])
        elif dll in imports:
            calls = len(imports[dll])
        libs[idx] += calls
    tot = libs.sum()
    return ( libs / tot ) if tot > 0 else libs # normalize

PE Section特征

对 PESection的一些信息进行编码，例如包含代码的部分与包含数据的部分的数量、标记为可执行的部分、每个部分的平均香农熵以及它们的大小与虚拟大小的平均比率 - 这些数据点将告诉模型 PE 是否以及如何被打包/压缩/混淆：

def encode_sections(pe):
    sections = [{ 
        'characteristics': ','.join(map(str, s.characteristics_lists)),
        'entropy': s.entropy,
        'name': s.name,
        'size': s.size,
        'vsize': s.virtual_size } for s in pe.sections]

    num_sections = len(sections)
    max_entropy= max(收缩for s in sections]) if num_sections else 0.0
    max_size= max(收缩for s in sections]) if num_sections else 0.0
    min_vsize= min(收缩for s in sections]) if num_sections else 0.0
    norm_size= (max_size / min_vsize) if min_vsize > 0 else 0.0

    return [ 
        # code_sections_ratio
        (len(收缩]) / num_sections) if num_sections else 0,
        # pec_sections_ratio
        (len(收缩]) / num_sections) if num_sections else 0,
        # sections_avg_entropy
        ((sum(收缩for s in sections]) / num_sections) / max_entropy) if max_entropy > 0 else 0.0,
        # sections_vsize_avg_ratio
        ((sum(收缩/ s['vsize'] for s in sections]) / num_sections) / norm_size) if norm_size > 0 else 0.0,
        ]

特征数据汇总处理

将上述一个所有特征属性数据汇总到一起，代表一个PE文件的全部特征向量数据

def encode_pe(filepath):
    log.debug("encoding %s ...", filepath)

    if hasattr(filepath, 'read'):
        raw = filepath.read()

    else:
        with open(filepath, 'rb') as fp:
            raw = fp.read()

    sz= len(raw)
    pe= lief.PE.parse(list(raw))
    ep_bytes = [0] * 64
    try:
        ep_offset = pe.entrypoint - pe.optional_header.imagebase
        ep_bytes= [int(b) for b in raw[ep_offset:ep_offset+64]]
    except Exception as e:
        log.warning("can't get entrypoint bytes from %s: %s", filepath, e)

    v = np.concatenate([ 
        encode_properties(pe),
        encode_entrypoint(ep_bytes),
        encode_histogram(raw),
        encode_libraries(pe),
        [ min(sz, pe.virtual_size) / max(sz, pe.virtual_size)],
        encode_sections(pe)
        ])

    return v

通过自定义ergo 项目中prepare.py文件的prepare_input将恶意文件输入并调用encode_pe函数获取到对应文件的特征向量数据

def prepare_input(x, is_encoding = False):
    # file upload
    if isinstance(x, werkzeug.datastructures.FileStorage):
        return encoder.encode_pe(x)
    # file path
    elif os.path.isfile(x) :
        return encoder.encode_pe(x)
    # raw vector
    else:
        return x.split(',')

执行后我们将获取对应特征工程的数据

ergo encode /path/to/ergo-pe-sec /path/to/dataset --output /path/to/dataset.csv

模型训练

将使用人工神经网络（ANN）的计算结构、在Adam 优化算法对其进行训练

人工神经网络（ANN）是一种模仿生物神经网络的计算模型。它通过大量的节点（或“神经元”）连接起来，形成一个多层的网络结构，可以用于各种任务，如分类、回归和模式识别。

ANN的基本组成部分包括：

• 输入层：接收输入数据。
• 隐藏层：处理输入数据，通常有多层，负责学习特征和模式。
• 输出层：产生模型的最终输出。

在训练过程中，ANN使用反向传播算法来调整权重，以最小化预测误差。常见的激活函数有ReLU、Sigmoid和Tanh等。

假设我们的数据集中的数据点之间存在数值相关性，我们对此并不了解，但如果知道，我们将能够将该数据集划分为输出类。我们所做的是要求这个黑盒获取数据集并通过迭代调整其内部参数来近似此类函数

在创建的项目model.py文件中可以发现我们的 ANN 的实现，有两个隐藏层，每个层有 70 个神经元，ReLU作为激活函数，在训练期间有 30% 的dropout（正则化技术，主要用于防止神经网络过拟合。其基本思想是在训练过程中随机“丢弃”一部分神经元，使得网络在每个训练批次中只使用一部分神经元进行前向传播和反向传播）

n_inputs = 486

return Sequential([
    Dense(70, input_shape=(n_inputs,), activation='relu'),
    Dropout(0.3),
    Dense(70, activation='relu'),
    Dropout(0.3),
    Dense(2, activation='softmax')
])

使用ergo开始模型训练过程，根据 CSV 文件中向量的总量，此过程可能需要几分钟、几小时甚至几天的时间。如果您的机器上有 GPU，Ergo 将自动使用它们而不是 CPU 核心，以显著加快训练速度

ergo train /ergo-pe-sec --dataset /dataset.csv

模型评估

等待模型训练完成后，可以使用ergo view来查看模型性能统计数据

ergo view /ergo-pe-sec

将显示训练历史记录，我们可以验证模型准确度确实随着时间的推移而增加（在我们的例子中，在第 30 个时期左右达到了 97％ 的准确度），以及ROC 曲线，它告诉我们模型可以有效地区分恶意与否

此外，还将显示每个训练集、验证集和测试集的混淆矩阵。左上角的对角线值（深红色）表示正确预测的数量，而其他值（粉红色）表示错误预测的数量（我们的模型在约 30000 个样本的测试集上有 1.4% 的误报率）：

工程化验证

上述使用ergo已经写好了特征化工程、模型训练等脚本，现在只需要将该项目跑起来即可；按照ergo工具的使用方式，先删除掉临时文件

ergo clean /ergo-pe-sec

加载模型并将其用作 API：

ergo serve /ergo-pe-sec --classes "safe, malicious"

在本地向该检测服务上传一个恶意文件验证下效果

curl -F "x=@/evil.exe" "http://localhost:8080/"

可以看到很明显已经检测出来为恶意文件的概率为0.999999999999

本文原文始发于微信公众号（暴暴的皮卡丘）