import os

import time

from PIL import Image

from tqdm import tqdm

import numpy as np

import pandas as pd

import pickle

from sklearn.model_selection import train_test_split

import tensorflow as tf

from tensorflow.keras import layers, models

from tensorflow.keras.preprocessing.image import ImageDataGenerator

# 超参数

TRAIN_SIZE = 0.8

VAL_SIZE = 0.1

TEST_SIZE = 0.1

SEED = 4396

LENGTH = 512

WIDTH, HEIGHT = 32, 16

BATCH_SIZE = 16

EPOCH = 300

SHUFFLE = False

CLASSES = 2

LR = 1e-4

datapath = "/home/jovyan/histogram"

with open("/home/datacon/malware/XXX/black.txt", 'r') as f:

black_list = f.read().strip().split()

with open("/home/datacon/malware/XXX/white.txt", 'r') as f:

white_list = f.read().strip().split()

black_path = [os.path.join(datapath, sp) for sp in black_list]

white_path = [os.path.join(datapath, sp) for sp in white_list]

raw_feature, raw_labels = [], []

with tqdm(total=11647, ncols=80, desc="histogram") as pbar:

for fp in black_path:

with open(fp+'.txt', 'r') as f:

feature = f.readlines()

feature = [float(his.strip()) for his in feature]

raw_feature.append(feature)

raw_labels.append(1)

pbar.update(1)

for fp in white_path:

with open(fp+'.txt', 'r') as f:

feature = f.readlines()

feature = [float(his.strip()) for his in feature]

raw_feature.append(feature)

raw_labels.append(0)

pbar.update(1)

# 打乱顺序

np.random.seed(SEED)

tf.random.set_seed(SEED)

features, labels = np.array(raw_feature, dtype=np.float32), np.array(raw_labels, dtype=np.int32)

index = list(range(len(labels)))

np.random.shuffle(index)

features = features[index]

labels = labels[index]

# 划分数据集

train_features, test_features, train_label, test_label = train_test_split(

features,

labels,

test_size=TEST_SIZE,

stratify=labels,

random_state=SEED)

train_features, valid_features, train_label, valid_label = train_test_split(

train_features,

train_label,

test_size=VAL_SIZE,

stratify=train_label,

random_state=SEED)

# 加载dataset

train_ds = tf.data.Dataset.from_tensor_slices((train_features, train_label)) \

.batch(BATCH_SIZE) \

.prefetch(buffer_size = tf.data.experimental.AUTOTUNE)

valid_ds = tf.data.Dataset.from_tensor_slices((valid_features, valid_label)) \

.batch(BATCH_SIZE) \

.prefetch(buffer_size = tf.data.experimental.AUTOTUNE)

test_ds = tf.data.Dataset.from_tensor_slices((test_features, test_label)) \

.batch(BATCH_SIZE) \

.prefetch(buffer_size = tf.data.experimental.AUTOTUNE)

# 模型

inputs = layers.Input(shape=(LENGTH, 1), dtype='float32')

re_inputs = tf.reshape(inputs, [-1, WIDTH, HEIGHT, 1])

Conv_1 = layers.Conv2D(60, (2, 2), padding='same', activation='relu')(re_inputs)

pool_1 = layers.MaxPooling2D()(Conv_1)

Conv_2 = layers.Conv2D(200, (2, 2), padding='same', activation='relu')(pool_1)

pool_2 = layers.MaxPooling2D()(Conv_2)

Flat = layers.Flatten()(pool_2)

Dense_1 = layers.Dense(500, activation='relu')(Flat)

dropout = layers.Dropout(0.2)(Dense_1)

# Dense_2 = layers.Dense(50, activation='relu')(dropout)

outputs = layers.Dense(1, activation='sigmoid')(Dense_1)

model = models.Model(inputs=inputs, outputs=outputs)

model.compile(optimizer=tf.keras.optimizers.Nadam(LR),

loss='binary_crossentropy',

metrics=['accuracy'])

model.fit(train_ds,

validation_data=valid_ds,

# class_weight=class_weight_dict,

epochs=EPOCH,

workers=4,

callbacks=[tf.keras.callbacks.EarlyStopping(monitor='val_loss', patience=6, min_delta=1e-4, mode='min'),

tf.keras.callbacks.ReduceLROnPlateau(monitor='val_loss', patience=4, factor=0.5, verbose=0)])

predict = model.evaluate(test_ds)

print(predict)

model.save('./models/histogram_{0:.2f}.h5'.format(predict[1]), save_format="tf")