ロード・トゥ・ザ・ホワイトハッカー

ホワイトハッカーはじめました

機械学習によるマルウェア判定（PEファイル）

Python
PEファイルのおさらい

PEファイルについて、以下を参照。

chikuwamarux.hatenablog.com

 

マルウェアの検体を取得する方法は以下を参照。

chikuwamarux.hatenablog.com

PEファイルの情報は、pefileというpythonのライブラリで取得できる。

問題のない正常なファイルはPCにいくらでもあるので、ディレクトリを再帰的に処理して、情報を取得すれば良い。ただ、それを上手に正規化されたデータに落とす方法を見つけることが出来なかった。

そういうデータセットを作成してくれている人がいるようで、Prateek Lalwaniという人が公開している、とのことですが、そのgithubはクローズされています。

ただ、他の人のgithubには公開されているので、そのデータを借用します。

 

以下２つのサイトで、MalwareData.csvのデータが取得できます。

github.com

上記は、そもそも、こちらの本を参考にさせてもらっていて、そのgithubサイトです。

 

github.com

どちらのサイトも、ファイルだけでなく、Pythonでの機械学習のコードがあります。

Kaggleにもデータがありますが、上記と同じもののようです。

www.kaggle.com

 

データについて

データの取り込み

import pandas as pd
MalwareDataset = pd.read_csv('MalwareData.csv', sep='|')

上記によりデータセットへ取り込みます。データは138047 rows × 57 columnsという内容。

MalwareDataset.columns

から、カラムを出力してみると、以下の57カラムであることが分かる。

Index(['Name', 'md5', 'Machine', 'SizeOfOptionalHeader', 'Characteristics', 'MajorLinkerVersion', 'MinorLinkerVersion', 'SizeOfCode', 'SizeOfInitializedData', 'SizeOfUninitializedData', 'AddressOfEntryPoint', 'BaseOfCode', 'BaseOfData', 'ImageBase', 'SectionAlignment', 'FileAlignment', 'MajorOperatingSystemVersion', 'MinorOperatingSystemVersion', 'MajorImageVersion', 'MinorImageVersion', 'MajorSubsystemVersion', 'MinorSubsystemVersion', 'SizeOfImage', 'SizeOfHeaders', 'CheckSum', 'Subsystem', 'DllCharacteristics', 'SizeOfStackReserve', 'SizeOfStackCommit', 'SizeOfHeapReserve', 'SizeOfHeapCommit', 'LoaderFlags', 'NumberOfRvaAndSizes', 'SectionsNb', 'SectionsMeanEntropy', 'SectionsMinEntropy', 'SectionsMaxEntropy', 'SectionsMeanRawsize', 'SectionsMinRawsize', 'SectionMaxRawsize', 'SectionsMeanVirtualsize', 'SectionsMinVirtualsize', 'SectionMaxVirtualsize', 'ImportsNbDLL', 'ImportsNb', 'ImportsNbOrdinal', 'ExportNb', 'ResourcesNb', 'ResourcesMeanEntropy', 'ResourcesMinEntropy', 'ResourcesMaxEntropy', 'ResourcesMeanSize', 'ResourcesMinSize', 'ResourcesMaxSize', 'LoadConfigurationSize', 'VersionInformationSize', 'legitimate'], dtype='object')

「legitimate」カラムがマルウェアかどうかの判定になっており、

という分類がされている。

MalwareDataset['legitimate'].value_counts()

でデータをカウントしてみると、

0：96724

1：41323

ということで、マルウェア多めのデータとなっている。

ここまでデータがあれば、あとは機械学習の分類アルゴリズムを適用できる。ただ、上記サンプルでは、機械学習の前に、ExtraTreeを使って、分類に寄与するパラメーターを絞って学習させることで、分類精度を上げることが可能になる。

つまり、どのPEヘッダーの情報がマルウェアの分類に影響するのか把握できる。

 

以下のコードにより、

(138047, 57) (138047, 13)　が出力され、13個の特徴量が分類に大きな影響を与えることがわかる。

from sklearn.ensemble import ExtraTreesClassifier      
from sklearn.feature_selection import SelectFromModel  

# 不要なカラムを削除
X = MalwareDataset.drop(['Name','md5','legitimate'],axis=1)

# データセットのラベル列のみを抽出してyに代入
y = MalwareDataset['legitimate'].values 

# ExtraTreesClassifierを使用
extratrees = ExtraTreesClassifier().fit(X, y)

# SelectFromModelを使用して、
# ExtraTreesClassifierによる分類結果に寄与した重要度の大きい特徴量のみを抽出
selection =  SelectFromModel(extratrees,prefit=True)

# 重要度の大きい特徴量のカラム名を取得
feature_idx = selection.get_support()
feature_name = X.columns[feature_idx]

new_data = selection.transform(X)
new_data = pd.DataFrame(new_data)
new_data.columns = feature_name
#Checking the shape of old as well as new data
print(MalwareDataset.shape,new_data.shape)

 

重要な特徴量を出力するコード

features = new_data.shape[1]

# 重要度をリストで抽出
importances = extratrees.feature_importances_

# 重要度を高い順にソート
indices = np.argsort(importances)[::-1]

# 重要度の高い順に、特徴量の名前と重要度を出力
for i in range(features):
    print("%d"%(i+1),MalwareDataset.columns[2+indices[i]],importances[indices[i]])

上記コードにより、以下が出力される。

1 DllCharacteristics 0.15836263233530415
2 Machine 0.12007604719789777
3 Characteristics 0.10400284183445602
4 SectionsMaxEntropy 0.06542977624819674
5 Subsystem 0.059400268661077955
6 VersionInformationSize 0.057993969120311704
7 ImageBase 0.05571243209223071
8 MajorSubsystemVersion 0.047525009638160316
9 ResourcesMaxEntropy 0.043357550686231955
10 SizeOfOptionalHeader 0.029203678693686497
11 MajorOperatingSystemVersion 0.026403782923872843
12 SectionsMinEntropy 0.023331472771475084
13 ResourcesMinEntropy 0.0202876288981414

これは、ランダムに実行されるので、実行のたびに結果が異なる。

上記では、「DllCharacteristics」が特徴量として重要、ということになる。

 

あとは、以下のサイトからコピペさせてもらい、機械学習を実行。

github.com

 

optunaによる最適なパラメータの探索

from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score 
from sklearn.model_selection import train_test_split
import numpy as np
import optuna
from sklearn.model_selection import cross_validate

# データセットを訓練用とテスト用に分割
X_train, X_test, y_train, y_test = train_test_split(
    new_data, y, test_size=0.2, shuffle=True, random_state=101
    )

# RandomForestClassifierのハイパーパラメータ探索用のクラスを設定
class Objective_RF:
    def __init__(self, X, y):
        self.X = X
        self.y = y

    def __call__(self, trial):
        # 探索対象のパラメータの設定
        criterion = trial.suggest_categorical("criterion",
                                              ["gini", "entropy"])
        bootstrap = trial.suggest_categorical('bootstrap',
                                              ['True','False'])
        max_features = trial.suggest_categorical('max_features',
                                            ['auto', 'sqrt','log2'])
        min_samples_split = trial.suggest_int('min_samples_split',
                                              2, 5)
        min_samples_leaf = trial.suggest_int('min_samples_leaf',
                                             1,10)

        model = RandomForestClassifier(
            criterion = criterion,
            bootstrap = bootstrap,
            max_features = max_features,
            min_samples_split = min_samples_split,
            min_samples_leaf = min_samples_leaf
        )

        # 交差検証しながらベストのパラメータ探索を行う
        scores = cross_validate(model,
                                X=self.X,
                                y=self.y,
                                cv=5,
                                n_jobs=-1)
        
        # 5分割で交差検証した正解率の平均値を返す
        return scores['test_score'].mean()

# 探索の対象クラスを設定
objective = Objective_RF(X_train, y_train)
study = optuna.create_study()
# 最大で3分間探索を実行
study.optimize(objective, timeout=180)
# ベストのパラメータの出力
print('params:', study.best_params)
params: {'criterion': 'gini', 'bootstrap': 'False', 'max_features': 'auto', 'min_samples_split': 5, 'min_samples_leaf': 10}

ベストパラメータで学習
from sklearn.metrics import confusion_matrix
from sklearn.metrics import accuracy_score

# optunaの探索結果として得られたベストのパラメータを設定
model = RandomForestClassifier(
    criterion = study.best_params['criterion'],
    bootstrap = study.best_params['bootstrap'],
    max_features = study.best_params['max_features'],
    min_samples_split = study.best_params['min_samples_split'],
    min_samples_leaf = study.best_params['min_samples_leaf']
)

# モデルの訓練
model.fit(X_train, y_train)

# テスト用のデータを使用して予測
pred = model.predict(X_test)

# 予測結果とテスト用のデータを使って正解率と、混同行列を出力
print("Accuracy: {:.5f} %".format(100 * accuracy_score(y_test, pred)))
print(confusion_matrix(y_test, pred))
Accuracy: 99.08367 %
[[19338   128]
 [  125  8019]]

ランダムフォレストでの特徴量の表示
%matplotlib inline

feat_importances = pd.Series(
    model.feature_importances_,
    index=new_data.columns).sort_values(ascending=True)

feat_importances.plot(kind='barh')

ここでは、「ImaegeBase」が重要なパラメータ、ということになっており、ExtraTreeの結果「DllChracteristics」と異なる。

ちなみに、ExtraTreeによるパラメータの削減をせずに、ランダムフォレストを実行した場合、Accuracy: 99.13799 %というスコアが出て、こちらの方が結果が良くなる。

これは過学習のせい、なのかもしれないが、その辺りの調査は必要。