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

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

ディープラーニングによるマルウェア検出(実行編)

前回のおさらい

以下の記事で、EXEファイルを画像化することをやった。

chikuwamarux.hatenablog.com

これにより、マルウェアを画像化して、AIに「似ている」ということを判断させることが可能になる。

 

データは以下を利用させてもらう。

www.dropbox.com

このデータは、

マルウェア名のフォルダに、そのマルウェアの画像が含まれている。

画像データの数値化処理

テレビと同じような原理で、映像を碁盤目に分けて、その1マス(画素)の色を数値化する。カラーの場合はRGBなので、(12, 34, 56)という3次元で表現される。

白黒の場合は、ゼロかイチかで表現できる。

例えば、2x2マスの白黒画像を数値化すると、以下のようになる。

2x2マスの白黒画像=[1,0,0,1]と表現できる。

実際は、もっと大きい画素で、ゼロとイチではなく、黒が占める割合で表現される。

画像ごとに処理を行いデータを作ることになる。

こんなコードを書けば、一発でやってくれてしまう。

from keras.preprocessing.image import ImageDataGenerator

#Generating DataSet
path_root = "malimg_paper_dataset_imgs\\"
batches = ImageDataGenerator().flow_from_directory(directory=path_root, target_size=(64,64), batch_size=10000)
imgs, labels = next(batches)

target_sizeは画像のサイズで64x64の画像に縮小している。

labelsにマルウェアの名称が入るが、それはフォルダ名をセットしてくれる。

ディープラーニングでの処理

上記データをもとにディープラーニングによるマルウェア検出器の作成。

 

データの分割

from sklearn.model_selection import train_test_split
X_train, X_test, y_train, y_test = train_test_split(imgs/255.,labels, test_size=0.3)

モデルの作成

import keras
from tensorflow.keras import Input
from tensorflow.keras.layers import BatchNormalization
from keras.models import Sequential,Model
from keras.layers import Dense, Dropout, Flatten
from keras.layers import Conv2D, MaxPooling2D

num_classes = 25

def malware_model():
    Malware_model = Sequential()
    Malware_model.add(Conv2D(30, kernel_size=(33),
                     activation='relu',
                     input_shape=(64,64,3)))

    Malware_model.add(MaxPooling2D(pool_size=(22)))
    Malware_model.add(Conv2D(15, (33), activation='relu'))
    Malware_model.add(MaxPooling2D(pool_size=(22)))
    Malware_model.add(Dropout(0.25))
    Malware_model.add(Flatten())
    Malware_model.add(Dense(128, activation='relu'))
    Malware_model.add(Dropout(0.5))
    Malware_model.add(Dense(50, activation='relu'))
    Malware_model.add(Dense(num_classes, activation='softmax'))
    Malware_model.compile(loss='categorical_crossentropy', optimizer = 'adam', metrics=['accuracy'])
    return Malware_model

Malware_model = malware_model()

Malware_model.summary()

2次元畳み込み層を使って、プーリング層を追加し、Flattenで一次元に変換しています。

訓練の実行

Malware_model.fit(X_train, y_train, validation_data=(X_test, y_test), epochs=10)

検証

scores = Malware_model.evaluate(X_test, y_test)

検証結果は、95.43%の精度が出ました。

まずまずの結果がでます。

まとめ

マルウェアを画像化したデータを使って、ディープラーニングによるマルチクラス分類を実装したら、95.43%の精度がでました。

 

参考サイト

殆ど、内容を端折って、こちらのコードを借用しております。一部、インポート文を修正しております。

towardsdatascience.com