-
Notifications
You must be signed in to change notification settings - Fork 13
/
DexRay.py
170 lines (139 loc) · 7.33 KB
/
DexRay.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
import numpy as np
import tensorflow as tf
import random as python_random
import tensorflow_addons as tfa
import tensorflow.keras as keras
from PIL import Image
random_seed = 123456
np.random.seed(random_seed)
python_random.seed(random_seed)
tf.random.set_seed(random_seed)
import os
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
import sys
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Conv1D, MaxPooling1D
from tensorflow.keras.layers import Flatten, Dense
import argparse
def parseargs():
parser = argparse.ArgumentParser()
parser.add_argument("-p", "--path", help="The path to the directory that contains malware and goodware image folders", type=str, required=True)
parser.add_argument("-d", "--dir", help="The name of the directory where to save the model", type=str, required=True)
parser.add_argument("-f", "--file", help="The name of the file where to save the results of the evaluation", type=str, required=True)
args = parser.parse_args()
return args
def get_label(file_path):
parts = tf.strings.split(file_path, os.path.sep)
if parts[-2] == 'goodware':
return [0]
else:
return [1]
def get_image(path_img):
image = np.asarray(Image.open(path_img))
image = tf.convert_to_tensor(image, dtype_hint=None, name=None)
return image
def get_shape(image):
return image.shape[0]
def decode_img(path_img):
image = tf.numpy_function(get_image, [path_img], tf.uint8)
shape = tf.numpy_function(get_shape, [image], tf.int64)
image = tf.reshape(image, [shape, 1, 1])
image = tf.image.convert_image_dtype(image, tf.float32)
image = tf.image.resize(image, [IMG_SIZE*IMG_SIZE, 1])
return tf.reshape(image, [IMG_SIZE*IMG_SIZE, 1])
def process_path(file_path):
label = get_label(file_path)
img = decode_img(file_path)
return img, label
def main(path_images, dir_name, file_name, CHANNELS, EPOCHS, BATCH_SIZE, IMG_SIZE, PATH_FILES, CLASS_NAMES):
recall_list, precision_list, accuracy_list, f1_list = [], [], [], []
model_architecture = Sequential()
model_architecture.add(Conv1D(filters=64, kernel_size=12, activation='relu', input_shape=(IMG_SIZE*IMG_SIZE, 1)))
model_architecture.add(MaxPooling1D(pool_size=12))
model_architecture.add(Conv1D(filters=128, kernel_size=12, activation='relu'))
model_architecture.add(MaxPooling1D(pool_size=12))
model_architecture.add(Flatten())
model_architecture.add(Dense(64, activation='sigmoid'))
model_architecture.add(Dense(1, activation='sigmoid'))
file_results = open(file_name, "w")
file_results.write("Scores of the performance evaluation are: Accuracy, Precision, Recall, F1-score\n")
for i in range(1, 11):
file_results.write("Run: %d \n" % i)
print("Run: %d" % i)
with open(os.path.join(PATH_FILES, "train"+str(i)+".txt")) as f:
train_hashes = f.read().splitlines()
train_imgs = [os.path.join(path_images, image_hash) for image_hash in train_hashes]
f.close()
with open(os.path.join(PATH_FILES, "valid"+str(i)+".txt")) as f:
valid_hashes = f.read().splitlines()
valid_imgs = [os.path.join(path_images, image_hash) for image_hash in valid_hashes]
f.close()
with open(os.path.join(PATH_FILES, "test"+str(i)+".txt")) as f:
test_hashes = f.read().splitlines()
test_imgs = [os.path.join(path_images, image_hash) for image_hash in test_hashes]
f.close()
train_dataset = tf.data.Dataset.from_tensor_slices(train_imgs)
train_dataset = train_dataset.map(process_path, num_parallel_calls=tf.data.experimental.AUTOTUNE)
length_train = len(train_imgs)
batch_train = length_train//BATCH_SIZE
train_dataset = train_dataset.cache()
train_dataset = train_dataset.shuffle(buffer_size=length_train, seed = random_seed, reshuffle_each_iteration=False)
train_dataset = train_dataset.batch(batch_train)
train_dataset = train_dataset.prefetch(tf.data.experimental.AUTOTUNE)
valid_dataset = tf.data.Dataset.from_tensor_slices(valid_imgs)
valid_dataset = valid_dataset.map(process_path, num_parallel_calls=tf.data.experimental.AUTOTUNE)
length_valid = len(valid_imgs)
batch_valid = length_valid//BATCH_SIZE
valid_dataset = valid_dataset.cache()
valid_dataset = valid_dataset.shuffle(buffer_size=length_valid, seed = random_seed, reshuffle_each_iteration=False)
valid_dataset = valid_dataset.batch(batch_valid)
valid_dataset = valid_dataset.prefetch(tf.data.experimental.AUTOTUNE)
test_dataset = tf.data.Dataset.from_tensor_slices(test_imgs)
test_dataset = test_dataset.map(process_path, num_parallel_calls=tf.data.experimental.AUTOTUNE)
length_test = len(test_imgs)
batch_test = length_test//BATCH_SIZE
test_dataset = test_dataset.cache()
test_dataset = test_dataset.shuffle(buffer_size=length_test, seed = random_seed, reshuffle_each_iteration=False)
test_dataset = test_dataset.batch(batch_test)
test_dataset = test_dataset.prefetch(tf.data.experimental.AUTOTUNE)
model = keras.models.clone_model(model_architecture)
model.compile(optimizer='adam',
loss=tf.keras.losses.BinaryCrossentropy(),
metrics=['accuracy',
tf.keras.metrics.Precision(),
tf.keras.metrics.Recall(),
tfa.metrics.F1Score(num_classes=2, average="micro", threshold=0.5)])
es_callback = tf.keras.callbacks.EarlyStopping(monitor='val_accuracy', patience=50, restore_best_weights=True)
cp_callback = tf.keras.callbacks.ModelCheckpoint(filepath=os.path.join(dir_name, 'cp'+str(i)),
save_weights_only=True,
monitor='val_accuracy',
mode='max',
save_best_only=True)
path_save_model = os.path.join(dir_name, 'model'+str(i))
model.fit(train_dataset, shuffle=True, validation_data = valid_dataset, epochs=EPOCHS, callbacks=[es_callback, cp_callback], verbose=2)
model.save(path_save_model)
print("Evaluate the model")
evaluation_scores = model.evaluate(test_dataset, verbose=2)
file_results.write("%s \n" % evaluation_scores[1:])
file_results.write("#"*50+"\n")
accuracy_list.append(evaluation_scores[1])
precision_list.append(evaluation_scores[2])
recall_list.append(evaluation_scores[3])
f1_list.append(evaluation_scores[4])
file_results.write("Average scores: %f %f %f %f" % (np.mean(accuracy_list),
np.mean(precision_list),
np.mean(recall_list),
np.mean(f1_list)))
file_results.close()
if __name__ == "__main__":
args = parseargs()
path_images = args.path
dir_name = args.dir
file_name = args.file
CHANNELS = 1
EPOCHS = 200
BATCH_SIZE = 500
IMG_SIZE = 128
PATH_FILES = "data_splits"
CLASS_NAMES = ['goodware', 'malware']
main(path_images, dir_name, file_name, CHANNELS, EPOCHS, BATCH_SIZE, IMG_SIZE, PATH_FILES, CLASS_NAMES)