使用机器学习和深度学习对PE进行二分类和多分类

怼烎@ 2022-09-03 00:15 13阅读 0赞

### 文章目录 ###

*  1. 前言
 *  2. 安装库
 *   *  2.1 安装lief
     *  2.2 安装ember
 *  3. 数据预处理
 *  4. 训练模型
 *  5. 预测新数据
 *   *  5.1 二分类预测
     *  5.2 多分类预测
 *  6. 代码下载

# 1. 前言 #

根据调研发现sophos-ai分享了两千万的PE数据集，链接为[https://github.com/sophos-ai/SOREL-20M][https_github.com_sophos-ai_SOREL-20M]。由于数据量规模大，所以就可以拿来训练较为有效的模型，不仅可以做对PE进行正常或者恶意的二分类，而且还可以进行更小粒度的恶意分类(如木马、蠕虫、勒索病毒等等)。

在静态条件下对PE进行分类(不把PE放到沙箱中执行)，主要有两种方法：

1.  提取特征后进行分类，例如使用`LIEF`进行特征提取后再进行分类。
2.  将PE文件看作是灰度图，然后将其视作是图像进行分类。

本文主要讲究的是方法一，即**通过工具进行特征提取后进行分类**。

# 2. 安装库 #

## 2.1 安装lief ##

pip install lief==0.11.4

源码链接：[https://github.com/lief-project/LIEF][https_github.com_lief-project_LIEF]

## 2.2 安装ember ##

pip install git+https://github.com/elastic/ember.git

源码链接：[https://github.com/elastic/ember][https_github.com_elastic_ember]

# 3. 数据预处理 #

假设每个文件名的格式为md5\_label。比如二分类，则label为0，则表示为白文件；label为1，则表示为恶意文件。如果是多分类，比如进行11分类：

1.  adware
2.  flooder
3.  ransomware
4.  dropper
5.  spyware
6.  packed
7.  crypto\_miner
8.  fle infector
9.  installer
10. worm
11. downloader

导包：

import sys
    import os
    import glob
    import ember
    import json

将结果保存到json文件中：

pe_extractor = ember.features.PEFeatureExtractor()
    json_content_list = []
    
    for i in glob.iglob('/home/learn/PE/data/*'):
        label = os.path.basename(i).split('_')[1]
        if label == 'nan':
            continue
        with open(i, 'rb') as f:
            json_content = pe_extractor.raw_features(f.read())
            json_content['label'] = int(float(label))
            json_content_list.append(json_content)
    
    with open('/home/learn/PE/result/processed_data.json', 'w') as f:
        f.write('\n'.join(json.dumps(i) for i in json_content_list))

读取json文件，并将特征和标签保存到dat类型文件中：

data_dir = '/home/learn/PE/data/result/'
    extractor = ember.features.PEFeatureExtractor()
    print("Vectorizing training set")
    X_path = os.path.join(data_dir, "X_train.dat")
    y_path = os.path.join(data_dir, "y_train.dat")
    raw_feature_paths = [os.path.join(data_dir, "processed_data.json")]
    nrows = sum([1 for fp in raw_feature_paths for line in open(fp)])
    ember.vectorize_subset(X_path, y_path, raw_feature_paths, extractor, nrows)
    
    X_train, y_train = ember.read_vectorized_features(data_dir, "train", 2)

# 4. 训练模型 #

比如使用lightgbm来训练模型，代码为：

import lightgbm as lgb
    params = {   "boosting": "gbdt",
                "objective": "binary",
                "num_iterations": 1000,
                "learning_rate": 0.05,
                "num_leaves": 100,
                "max_depth": 15,
                "min_data_in_leaf": 20,
                "feature_fraction": 0.5,
    }
    
    train_rows = (y_train != -1)
    
    # Train
    lgbm_dataset = lgb.Dataset(X_train[train_rows], y_train[train_rows].astype('int32'))
    lgbm_model = lgb.train(params, lgbm_dataset)
    
    
    lgbm_model.save_model('lightgbm.model')

# 5. 预测新数据 #

## 5.1 二分类预测 ##

import lightgbm as lgb
    import ember
    
    binary_path = 'putty.exe'
    file_data = open(binary_path, "rb").read()
    lgbm_model = lgb.Booster(model_file='lightgbm.model')
    score = ember.predict_sample(lgbm_model, file_data)
    print(score)

## 5.2 多分类预测 ##

import ember
    import torch
    from torch import nn
    import torch.nn.functional as F
    
    class PENetwork(nn.Module):
        """ This is a simple network loosely based on the one used in ALOHA: Auxiliary Loss Optimization for Hypothesis Augmentation (https://arxiv.org/abs/1903.05700) Note that it uses fewer (and smaller) layers, as well as a single layer for all tag predictions, performance will suffer accordingly. """
        def __init__(self, use_malware=True, use_counts=False, use_tags=True, n_tags=11, feature_dimension=2381, layer_sizes = None):
            self.use_malware=use_malware
            self.use_counts=use_counts
            self.use_tags=use_tags
            self.n_tags = n_tags
    
            if self.use_tags and self.n_tags == None:
                raise ValueError("n_tags was None but we're trying to predict tags. Please include n_tags")
            super(PENetwork,self).__init__()
            p = 0.05
            layers = []
            if layer_sizes is None:layer_sizes=[512,512,128]
            for i,ls in enumerate(layer_sizes):
                if i == 0:
                    layers.append(nn.Linear(feature_dimension,ls))
                else:
                    layers.append(nn.Linear(layer_sizes[i-1],ls))
                layers.append(nn.LayerNorm(ls))
                layers.append(nn.ELU())
                layers.append(nn.Dropout(p))
            self.model_base = nn.Sequential(*tuple(layers))
            self.malware_head = nn.Sequential(nn.Linear(layer_sizes[-1], 1),
                                              nn.Sigmoid())
            self.count_head = nn.Linear(layer_sizes[-1], 1)
            self.sigmoid = nn.Sigmoid()
            self.tag_head = nn.Sequential(nn.Linear(layer_sizes[-1],64),
                                            nn.ELU(), 
                                            nn.Linear(64,64),
                                            nn.ELU(),
                                            nn.Linear(64,n_tags),
                                            nn.Sigmoid())
    
        def forward(self,data):
            rv = { }
            base_result = self.model_base.forward(data)
            if self.use_malware:
                rv['malware'] = self.malware_head(base_result)
            if self.use_counts:
                rv['count'] = self.count_head(base_result)
            if self.use_tags:
                rv['tags'] = self.tag_head(base_result)
            return rv
    
    pe_extractor = ember.features.PEFeatureExtractor()
    binary_path = 'malware.exe'
    file_data = open(binary_path, "rb").read()
    X_test = pe_extractor.feature_vector(file_data)
    
    with torch.no_grad():
        model = PENetwork()
        model.load_state_dict(torch.load('nn.pt'))
        model.eval()
        y_pred = model(torch.from_numpy(X_test))
        print(y_pred['tags'].argmax())

# 6. 代码下载 #

为了方便大家进行测试，提供的代码包含的是用保存后的模型对数据进行预测的代码，下载链接为：[https://download.csdn.net/download/herosunly/20721893][https_download.csdn.net_download_herosunly_20721893]。

[https_github.com_sophos-ai_SOREL-20M]: https://github.com/sophos-ai/SOREL-20M
[https_github.com_lief-project_LIEF]: https://github.com/lief-project/LIEF
[https_github.com_elastic_ember]: https://github.com/elastic/ember
[https_download.csdn.net_download_herosunly_20721893]: https://download.csdn.net/download/herosunly/20721893