东南大学软件安全与恶意代码分析大作业

import os import pickle from sklearn.ensemble import RandomForestClassifier import pefile def get_peheader_features(path): peheader_features = [] try: pe = pefile.PE(path) if hasattr(pe, 'OPTIONAL_HEADER'): cur = pe.OPTIONAL_HEADER for it in cur.__keys__: peheader_features += [cur.__getattribute__(it[0])] if(len(peheader_features) != 30): return [0] * 30 print(path) return peheader_features except pefile.PEFormatError: return [0] * 30 def train_detector(): os.chdir(r"D:\PyCharmProjects\MiningDetection\black_fixed") black_paths = os.listdir() os.chdir(r"D:\PyCharmProjects\MiningDetection\white_fixed") white_paths = os.listdir() for i in range(len(black_paths)): black_paths[i] = r"D:\PyCharmProjects\MiningDetection\black_fixed\\" + black_paths[i] for i in range(len(white_paths)): white_paths[i] = r"D:\PyCharmProjects\MiningDetection\white_fixed\\" + white_paths[i] print("Begin Training...") X = [get_peheader_features(path) for path in black_paths + white_paths] y = [1 for i in range(len(black_paths))] + [0 for i in range(len(white_paths))] classifier = RandomForestClassifier(n_estimators=500) classifier.fit(X, y) print("End Training...") os.chdir(r"D:\PyCharmProjects\MiningDetection") pickle.dump(classifier, open("train_data.pkl", "wb+")) train_detector()