####################################################
# [ 演習1 ] pefileを使ってファイルアクセスをしてみる
# pefile : PEファイルにアクセスするためのmodule
# `pip install pefile`
# exeファイルはPEヘッダに値を複数持っている
# DOS_HEADER/FILE_HEADER/OPTIONAL_HEADERを参照する
####################################################
import pefile

PATH = 'data/PE-samples/TestApp.exe'
pe = pefile.PE(PATH)
print("{0}".format(pe.DOS_HEADER))
print("{0}".format(pe.FILE_HEADER))
print("{0}".format(pe.OPTIONAL_HEADER))

if pe.FILE_HEADER.Characteristics & 0x0002:
    print('IMAGE_FILE_EXECUTABLE_IMAGE')
if pe.FILE_HEADER.Characteristics & 0x0100:
    print('IMAGE_FILE_32BIT_MACHINE')

####################################################
# [ 演習2 ] 正規なデータのPEヘッダの値を覗いてみる
# pandas : データ操作をしやすくするもの(配列など)
# data/PE-samples/以下に特徴抽出対象のファイルがある
# PEヘッダの値を特徴量として抽出する
####################################################
import glob
import pandas as pd

df = pd.DataFrame(columns = ['malware', 'VirtualAddress', 'ResourceSize', 'DebugSize', 'IATSize'])
PATH = 'data/PE-samples/*'
files = glob.glob(PATH)
for file in files:
    data = pefile.PE(file)
    VA = data.OPTIONAL_HEADER.DATA_DIRECTORY[1].VirtualAddress
    RS = data.OPTIONAL_HEADER.DATA_DIRECTORY[2].Size
    DS = data.OPTIONAL_HEADER.DATA_DIRECTORY[6].Size
    IATSize = data.OPTIONAL_HEADER.DATA_DIRECTORY[1].Size
    newdf = pd.DataFrame([[0, VA, RS, DS, IATSize]], columns = ['malware', 'VirtualAddress', 'ResoureSize', 'DebugSize', 'IATSize'])
    df = df.append(newdf, ignore_index = True)
df

####################################################
# [ 演習3 ] データの可視化
## malware.csv : マルウェアから抽出したデータファイル
## benign.csv : 正規ファイルから抽出したデータファイル
## マルウェアから抽出した値を赤色で、
## 正規データから抽出した値を青色で表し二次元で描画する
####################################################
import numpy as np
import matplotlib.pyplot as plt
from sklearn.decomposition import PCA

df = pd.read_csv('data/malware.csv')
df2 = pd.read_csv('data/benign.csv')

X_reduced = PCA(n_components = 2).fit_transform(np.array(df)[:, 2:])
Y_reduced = PCA(n_components = 2).fit_transform(np.array(df2)[:, 2:])

plt.scatter(X_reduced[:, 0], X_reduced[:, 1], c = 'red')
plt.scatter(Y_reduced[:, 0], Y_reduced[:, 1], c = 'blue')
plt.show()

####################################################
# [ 演習4 ] 教師ありデータの作成
## 次元削減をして一次元で描画する
####################################################
from sklearn.discriminant_analysis import LinearDiscriminantAnalysis as LDA

df = pd.read_csv('data/union.csv')
X_reduced = LDA(n_components = 1).fit_transform(np.array(df)[:, 2:], np.array(df)[:, 1])
plt.plot(X_reduced[:, 0], c = 'red')
plt.hlines(y = 0.2, xmin = 0, xmax = 355, linestyles = 'dashed')
plt.show()

####################################################
# [ 演習5 ] マルウェアかどうか予測する
####################################################
from sklearn.neighbors import KNeighborsClassifier

df = pd.read_csv('data/union.csv')
clf_k = KNeighborsClassifier(n_neighbors = 3)
clf_k.fit(np.array(df)[:, 2:], np.array(df)[:, 1])

test_df = pd.read_csv('data/test2.csv')
output = clf_k.predict(np.array(test_df)[:, 2:])

for i in output:
    if i == 0:
        print('not malware')
    elif i == 1:
        print('malware')

####################################################
# [ 演習6 ] 判定の精度を見てみる
## 56%なので割と精度が低い
####################################################
from sklearn.metrics import accuracy_score
print(accuracy_score(np.array(test_df)[:, 1], output)) # => 0.56

####################################################
# [ 演習7 ] 精度を上げるために抽出量のチューニングをする
## K近傍法で判定する
## Kの個数でチューニングを行う
## K=13で精度が 56% => 76% に向上していることがわかる
####################################################
for i in range(1, 20):
    clf_k = KNeighborsClassifier(n_neighbors = i)
    clf_k.fit(np.array(df)[:, 2:], np.array(df)[:, 1])
    output = clf_k.predict(np.array(test_df)[:, 2:])
    print("n_neighbors={0} {1}".format(i, accuracy_score(np.array(test_df)[:, 1], output)))