[scikit-learn]使用随机森林来识别恶意软件

最近发现树模型在数据分析中还是挺常见的，我这里把一些demo代码分享出来，方便大家进行后期的加工，用到自己的项目上，我这里使用的数据集是恶意软件的表格数据，下载地址为：

​​https://github.com/chihebchebbi/Mastering-Machine-Learning-for-Penetration-Testing/blob/master/Chapter03/MalwareData.csv.gz​​

下载完数据，解压放到data目录里面，然后使用随机森林的示例代码（rf_demo.py）为：

import pandas as pd

import sklearn
from sklearn.feature_selection import SelectFromModel
from sklearn.ensemble import ExtraTreesClassifier
from sklearn.metrics import confusion_matrix
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score

MalwareDataset = pd.read_csv('data/MalwareData.csv', sep='|')
Legit = MalwareDataset[0:41323].drop(['legitimate'], axis=1)
Malware = MalwareDataset[41323::].drop(['legitimate'], axis=1)

print('The Number of important features is %i \n' % Legit.shape[1])
Data = MalwareDataset.drop(['Name', 'md5', 'legitimate'], axis=1).values
Target = MalwareDataset['legitimate'].values
FeatSelect = sklearn.ensemble.ExtraTreesClassifier().fit(Data, Target)
Model = SelectFromModel(FeatSelect, prefit=True)
Data_new = Model.transform(Data)

Legit_Train, Legit_Test, Malware_Train, Malware_Test = train_test_split(Data_new, Target ,test_size=0.2)

clf = sklearn.ensemble.RandomForestClassifier(n_estimators=50)
clf.fit(Legit_Train, Malware_Train)

score = clf.score(Legit_Test, Malware_Test)
print("The score of Random Forest Algorithm is", score*100)
Result = clf.predict(Legit_Test)

acc = accuracy_score(Malware_Test, Result)
print("accuracy score is {}".format(acc))

参考文献

​​Chapter 3: Malware analysis using API Calls and PE headers​​