├── BenignGenerator.py ├── Malware_Detection.ipynb ├── README.md ├── build_image_data.py ├── cnn1.py └── ias_full.py /BenignGenerator.py: -------------------------------------------------------------------------------- 1 | import os 2 | import shutil 3 | def recursive_walk(folder): 4 | 5 | for folderName, subfolders, filenames in os.walk(folder): 6 | if subfolders: 7 | for subfolder in subfolders: 8 | recursive_walk(subfolder) 9 | #print('\nFolder: ' + folderName + '\n') 10 | for filename in filenames: 11 | if filename.endswith('.exe'): 12 | shutil.copy(folderName+"\\"+filename, dir_dst) 13 | #print(filename) 14 | unallowed = ['desktop.ini','WindowsApps'] 15 | l=os.listdir("C:\\Program Files\\") 16 | dir_src = ("C:\\Program Files\\") 17 | dir_dst = ("C:\\BenignFiles\\") 18 | for i in l: 19 | if i in unallowed: 20 | continue 21 | print('C:\\Program Files\\' +i) 22 | recursive_walk('C:\\Program Files\\'+i) 23 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Malware-Detection-using-Deep-Learning 2 | 3 | ## Installations: 4 | CUDA - 9.0 5 | CuDNN - 7.0.5 6 | Python - 3.5+ 7 | Anaconda - 3 8 | 9 | ## Python Libraries : 10 | -Tensorflow GPU enabled . 11 | -Windows version x86 exe(local). 12 | -Pillow. 13 | -Numpy. 14 | 15 | ## Steps: 16 | - If you don't have benign files, use the benignGenerator.py to crawl your local machines and collect .exe files, which you can use as nenign files in your dataset. 17 | - Update the paths in ias.py file according to the location of the files. Use absolute paths. 18 | - Make sure the folders in path2 and path3 (destination of the output images 256x256 and 32x32 respectively) are created beforehand. 19 | - Run the ias_full.py file. 20 | - Update the paths in cnn1.py and build_image_data.py files according to the location of the files. Use absolute paths. 21 | - Run build_image_data.py. Make sure that label.txt is in the same folder as build_image_data as well as the folders(named after the 22 | classes in label.txt) are in the same folder. 23 | - Run the build_image_data.py file. After running the file you should see "train-00000-of-00002.tfrecord" and 24 | "train-00000-of-00002.tfrecord" files created. 25 | - Run cnn1.py. This will take a long time to run. 26 | - Update the paths in cnn1.py and build_image_data.py files according to the location of the files. Use absolute paths. 27 | - Run build_image_data.py. Make sure that label.txt is in the same folder as build_image_data as well as the folders(named after the 28 | classes in label.txt) are in the same folder. 29 | - Run the build_image_data.py file. After running the file you should see "train-00000-of-00002.tfrecord" and 30 | "train-00000-of-00002.tfrecord" files created. 31 | - Run cnn1.py. This will take a long time to run. 32 | 33 | IEEE Paper link(base paper): http://ieeexplore.ieee.org/document/8190895/references?ctx=references 34 | 35 | The code in jupyter notebook uses fast.ai libraries to solve the same problem statement. The code has been written after following the first lecture of the fast.ai lecture. to understand the code go through the pet classification problem in the lecture. 36 | 37 | ### Note: The model implemented using fast.ai libraries gives much better results than the one implemented using tensorflow. 38 | -------------------------------------------------------------------------------- /build_image_data.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | from __future__ import division 3 | from __future__ import print_function 4 | 5 | from datetime import datetime 6 | import os 7 | import random 8 | import sys 9 | import threading 10 | 11 | 12 | import numpy as np 13 | import tensorflow as tf 14 | 15 | tf.app.flags.DEFINE_string('train_directory', './', 16 | 'Training data directory') 17 | tf.app.flags.DEFINE_string('validation_directory', './', 18 | 'Validation data directory') 19 | tf.app.flags.DEFINE_string('output_directory', './', 20 | 'Output data directory') 21 | 22 | tf.app.flags.DEFINE_integer('train_shards', 2, 23 | 'Number of shards in training TFRecord files.') 24 | tf.app.flags.DEFINE_integer('validation_shards', 0, 25 | 'Number of shards in validation TFRecord files.') 26 | 27 | tf.app.flags.DEFINE_integer('num_threads', 2, 28 | 'Number of threads to preprocess the images.') 29 | 30 | # The labels file contains a list of valid labels are held in this file. 31 | # Assumes that the file contains entries as such: 32 | # malware 33 | # benign 34 | # where each line corresponds to a label. We map each label contained in 35 | # the file to an integer corresponding to the line number starting from 0. 36 | tf.app.flags.DEFINE_string('labels_file', './label.txt', 'Labels file') 37 | 38 | 39 | FLAGS = tf.app.flags.FLAGS 40 | 41 | 42 | def _int64_feature(value): 43 | """Wrapper for inserting int64 features into Example proto.""" 44 | if not isinstance(value, list): 45 | value = [value] 46 | return tf.train.Feature(int64_list=tf.train.Int64List(value=value)) 47 | 48 | 49 | def _bytes_feature(value): 50 | """Wrapper for inserting bytes features into Example proto.""" 51 | return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value])) 52 | 53 | 54 | def _convert_to_example(filename, image_buffer, label, text, height, width): 55 | """Build an Example proto for an example. 56 | Args: 57 | filename: string, path to an image file, e.g., '/path/to/example.JPG' 58 | image_buffer: string, JPEG encoding of RGB image 59 | label: integer, identifier for the ground truth for the network 60 | text: string, unique human-readable, e.g. 'malware' 61 | height: integer, image height in pixels 62 | width: integer, image width in pixels 63 | Returns: 64 | Example proto 65 | """ 66 | 67 | colorspace = 'RGB' 68 | channels = 3 69 | image_format = 'JPEG' 70 | 71 | example = tf.train.Example(features=tf.train.Features(feature={ 72 | 'image/height': _int64_feature(height), 73 | 'image/width': _int64_feature(width), 74 | 'image/colorspace': _bytes_feature(tf.compat.as_bytes(colorspace)), 75 | 'image/channels': _int64_feature(channels), 76 | 'image/class/label': _int64_feature(label), 77 | 'image/class/text': _bytes_feature(tf.compat.as_bytes(text)), 78 | 'image/format': _bytes_feature(tf.compat.as_bytes(image_format)), 79 | 'image/filename': _bytes_feature(tf.compat.as_bytes(os.path.basename(filename))), 80 | 'image/encoded': _bytes_feature(tf.compat.as_bytes(image_buffer))})) 81 | return example 82 | 83 | 84 | class ImageCoder(object): 85 | """Helper class that provides TensorFlow image coding utilities.""" 86 | 87 | def __init__(self): 88 | # Create a single Session to run all image coding calls. 89 | self._sess = tf.Session() 90 | 91 | # Initializes function that converts PNG to JPEG data. 92 | self._png_data = tf.placeholder(dtype=tf.string) 93 | image = tf.image.decode_png(self._png_data, channels=3) 94 | self._png_to_jpeg = tf.image.encode_jpeg(image, format='rgb', quality=100) 95 | 96 | # Initializes function that decodes RGB JPEG data. 97 | self._decode_jpeg_data = tf.placeholder(dtype=tf.string) 98 | self._decode_jpeg = tf.image.decode_jpeg(self._decode_jpeg_data, channels=3) 99 | 100 | def png_to_jpeg(self, image_data): 101 | return self._sess.run(self._png_to_jpeg, 102 | feed_dict={self._png_data: image_data}) 103 | 104 | def decode_jpeg(self, image_data): 105 | image = self._sess.run(self._decode_jpeg, 106 | feed_dict={self._decode_jpeg_data: image_data}) 107 | assert len(image.shape) == 3 108 | assert image.shape[2] == 3 109 | return image 110 | 111 | 112 | def _is_png(filename): 113 | ## Determine if a file contains a PNG format image. 114 | return '.png' in filename 115 | 116 | 117 | def _process_image(filename, coder): 118 | ## Process a single image file. 119 | 120 | # Read the image file. 121 | with tf.gfile.FastGFile(filename, 'rb') as f: 122 | image_data = f.read() 123 | 124 | # Convert PNG to JPEG 125 | if _is_png(filename): 126 | print('Converting PNG to JPEG for %s' % filename) 127 | image_data = coder.png_to_jpeg(image_data) 128 | 129 | image = coder.decode_jpeg(image_data) ##instance of ImageCoder to provide TensorFlow image coding utils. 130 | 131 | # Check that image converted to RGB 132 | assert len(image.shape) == 3 133 | height = image.shape[0] 134 | width = image.shape[1] 135 | assert image.shape[2] == 3 136 | 137 | return image_data, height, width 138 | 139 | 140 | def _process_image_files_batch(coder, thread_index, ranges, name, filenames, 141 | texts, labels, num_shards): 142 | """Processes and saves list of images as TFRecord in 1 thread. 143 | Args: 144 | coder: instance of ImageCoder to provide TensorFlow image coding utils. 145 | thread_index: integer, unique batch to run index is within [0, len(ranges)). 146 | ranges: list of pairs of integers specifying ranges of each batches to 147 | analyze in parallel. 148 | name: string, unique identifier specifying the data set 149 | filenames: list of strings; each string is a path to an image file 150 | texts: list of strings; each string is human readable, e.g. 'malware' 151 | labels: list of integer; each integer identifies the ground truth 152 | num_shards: integer number of shards for this data set. 153 | """ 154 | # Each thread produces N shards where N = int(num_shards / num_threads). 155 | # For instance, if num_shards = 128, and the num_threads = 2, then the first 156 | # thread would produce shards [0, 64). 157 | num_threads = len(ranges) 158 | assert not num_shards % num_threads 159 | num_shards_per_batch = int(num_shards / num_threads) 160 | 161 | shard_ranges = np.linspace(ranges[thread_index][0], 162 | ranges[thread_index][1], 163 | num_shards_per_batch + 1).astype(int) 164 | num_files_in_thread = ranges[thread_index][1] - ranges[thread_index][0] 165 | 166 | counter = 0 167 | for s in range(num_shards_per_batch): 168 | # Generate a sharded version of the file name, e.g. 'train-00002-of-00010' 169 | shard = thread_index * num_shards_per_batch + s 170 | output_filename = '%s-%.5d-of-%.5d.tfrecord' % (name, shard, num_shards) 171 | output_file = os.path.join(FLAGS.output_directory, output_filename) 172 | writer = tf.python_io.TFRecordWriter(output_file) 173 | 174 | shard_counter = 0 175 | files_in_shard = np.arange(shard_ranges[s], shard_ranges[s + 1], dtype=int) 176 | for i in files_in_shard: 177 | filename = filenames[i] 178 | label = labels[i] 179 | text = texts[i] 180 | 181 | image_buffer, height, width = _process_image(filename, coder) 182 | 183 | example = _convert_to_example(filename, image_buffer, label, 184 | text, height, width) 185 | writer.write(example.SerializeToString()) 186 | shard_counter += 1 187 | counter += 1 188 | 189 | if not counter % 1000: 190 | print('%s [thread %d]: Processed %d of %d images in thread batch.' % 191 | (datetime.now(), thread_index, counter, num_files_in_thread)) 192 | sys.stdout.flush() 193 | 194 | writer.close() 195 | print('%s [thread %d]: Wrote %d images to %s' % 196 | (datetime.now(), thread_index, shard_counter, output_file)) 197 | sys.stdout.flush() 198 | shard_counter = 0 199 | print('%s [thread %d]: Wrote %d images to %d shards.' % 200 | (datetime.now(), thread_index, counter, num_files_in_thread)) 201 | sys.stdout.flush() 202 | 203 | 204 | def _process_image_files(name, filenames, texts, labels, num_shards): 205 | """Process and save list of images as TFRecord of Example protos. 206 | Args: 207 | name: string, unique identifier specifying the data set 208 | filenames: list of strings; each string is a path to an image file 209 | texts: list of strings; each string is human readable, e.g. 'malware' 210 | labels: list of integer; each integer identifies the ground truth 211 | num_shards: integer number of shards for this data set. 212 | """ 213 | assert len(filenames) == len(texts) 214 | assert len(filenames) == len(labels) 215 | 216 | # Break all images into batches with a [ranges[i][0], ranges[i][1]]. 217 | spacing = np.linspace(0, len(filenames), FLAGS.num_threads + 1).astype(np.int) 218 | ranges = [] 219 | for i in range(len(spacing) - 1): 220 | ranges.append([spacing[i], spacing[i+1]]) 221 | 222 | # Launch a thread for each batch. 223 | print('Launching %d threads for spacings: %s' % (FLAGS.num_threads, ranges)) 224 | sys.stdout.flush() 225 | 226 | # Create a mechanism for monitoring when all threads are finished 227 | coord = tf.train.Coordinator() 228 | 229 | # Create a generic TensorFlow-based utility for converting all image codings. 230 | coder = ImageCoder() 231 | 232 | threads = [] 233 | for thread_index in range(len(ranges)): 234 | args = (coder, thread_index, ranges, name, filenames, 235 | texts, labels, num_shards) 236 | t = threading.Thread(target=_process_image_files_batch, args=args) 237 | t.start() 238 | threads.append(t) 239 | 240 | # Wait for all the threads to terminate. 241 | coord.join(threads) 242 | print('%s: Finished writing all %d images in data set.' % 243 | (datetime.now(), len(filenames))) 244 | sys.stdout.flush() 245 | 246 | 247 | def _find_image_files(data_dir, labels_file): 248 | """Build a list of all images files and labels in the data set. 249 | Returns: 250 | filenames: list of strings; each string is a path to an image file. 251 | texts: list of strings; each string is the class, e.g. 'malware' 252 | labels: list of integer; each integer identifies the ground truth. 253 | """ 254 | print('Determining list of input files and labels from %s.' % data_dir) 255 | unique_labels = [l.strip() for l in tf.gfile.FastGFile( 256 | labels_file, 'r').readlines()] 257 | 258 | labels = [] 259 | filenames = [] 260 | texts = [] 261 | 262 | # Leave label index 0 empty as a background class. 263 | label_index = 1 264 | 265 | # Construct the list of JPEG files and labels. 266 | for text in unique_labels: 267 | jpeg_file_path = '%s/%s/*' % (data_dir, text) 268 | matching_files = tf.gfile.Glob(jpeg_file_path) 269 | 270 | labels.extend([label_index] * len(matching_files)) 271 | texts.extend([text] * len(matching_files)) 272 | filenames.extend(matching_files) 273 | 274 | if not label_index % 100: 275 | print('Finished finding files in %d of %d classes.' % ( 276 | label_index, len(labels))) 277 | label_index += 1 278 | 279 | # Shuffle the ordering of all image files 280 | shuffled_index = list(range(len(filenames))) 281 | random.seed(12345) 282 | random.shuffle(shuffled_index) 283 | 284 | filenames = [filenames[i] for i in shuffled_index] 285 | texts = [texts[i] for i in shuffled_index] 286 | labels = [labels[i] for i in shuffled_index] 287 | 288 | print('Found %d JPEG files across %d labels inside %s.' % 289 | (len(filenames), len(unique_labels), data_dir)) 290 | return filenames, texts, labels 291 | 292 | 293 | def _process_dataset(name, directory, num_shards, labels_file): 294 | # Process a complete data set and save it as a TFRecord. 295 | # Args: 296 | # name: string, unique identifier specifying the data set. 297 | # directory: string, root path to the data set. 298 | # num_shards: integer number of shards for this data set. 299 | # labels_file: string, path to the labels file. 300 | filenames, texts, labels = _find_image_files(directory, labels_file) 301 | _process_image_files(name, filenames, texts, labels, num_shards) 302 | 303 | 304 | def main(unused_argv): 305 | assert not FLAGS.train_shards % FLAGS.num_threads, ( 306 | 'Please make the FLAGS.num_threads commensurate with FLAGS.train_shards') 307 | assert not FLAGS.validation_shards % FLAGS.num_threads, ( 308 | 'Please make the FLAGS.num_threads commensurate with ' 309 | 'FLAGS.validation_shards') 310 | print('Saving results to %s' % FLAGS.output_directory) 311 | 312 | # Run it! 313 | _process_dataset('validation', FLAGS.validation_directory, 314 | FLAGS.validation_shards, FLAGS.labels_file) 315 | _process_dataset('train', FLAGS.train_directory, 316 | FLAGS.train_shards, FLAGS.labels_file) 317 | 318 | 319 | if __name__ == '__main__': 320 | tf.app.run() -------------------------------------------------------------------------------- /cnn1.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf # tensorflow module 2 | import numpy as np # numpy module 3 | import os # path join 4 | 5 | 6 | DATA_DIR = "C:/Users/admin/Project IAS/" 7 | TRAINING_SET_SIZE = 5812 8 | BATCH_SIZE = 10 9 | IMAGE_SIZE = 224 10 | 11 | 12 | def _int64_feature(value): 13 | return tf.train.Feature(int64_list=tf.train.Int64List(value=value)) 14 | 15 | def _bytes_feature(value): 16 | return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value])) 17 | 18 | # image object from protobuf 19 | class _image_object: 20 | def __init__(self): 21 | self.image = tf.Variable([], dtype = tf.string) 22 | self.height = tf.Variable([], dtype = tf.int64) 23 | self.width = tf.Variable([], dtype = tf.int64) 24 | self.filename = tf.Variable([], dtype = tf.string) 25 | self.label = tf.Variable([], dtype = tf.int32) 26 | 27 | ## extracting information and storing them in an image object. 28 | def read_and_decode(filename_queue): 29 | reader = tf.TFRecordReader() 30 | _, serialized_example = reader.read(filename_queue) 31 | features = tf.parse_single_example(serialized_example, features = { 32 | "image/encoded": tf.FixedLenFeature([], tf.string), 33 | "image/height": tf.FixedLenFeature([], tf.int64), 34 | "image/width": tf.FixedLenFeature([], tf.int64), 35 | "image/filename": tf.FixedLenFeature([], tf.string), 36 | "image/class/label": tf.FixedLenFeature([], tf.int64),}) 37 | image_encoded = features["image/encoded"] 38 | image_raw = tf.image.decode_jpeg(image_encoded, channels=3) 39 | image_object = _image_object() 40 | image_object.image = tf.image.resize_image_with_crop_or_pad(image_raw, IMAGE_SIZE, IMAGE_SIZE) 41 | image_object.height = features["image/height"] 42 | image_object.width = features["image/width"] 43 | image_object.filename = features["image/filename"] 44 | image_object.label = tf.cast(features["image/class/label"], tf.int64) 45 | return image_object 46 | 47 | ## read input images for training and testing from .tfrecord files. 48 | def malware_input(if_random = True, if_training = True): 49 | if(if_training): 50 | filenames = [os.path.join(DATA_DIR, "train-0000%d-of-00002.tfrecord" % i) for i in range(0, 1)] 51 | else: 52 | filenames = [os.path.join(DATA_DIR, "eval-0000%d-of-00002.tfrecord" % i) for i in range(0, 1)] ## changing eval to train here!! 53 | 54 | for f in filenames: 55 | if not tf.gfile.Exists(f): 56 | raise ValueError("Failed to find file: " + f) 57 | filename_queue = tf.train.string_input_producer(filenames) 58 | image_object = read_and_decode(filename_queue) 59 | image = tf.image.per_image_standardization(image_object.image) 60 | # image = image_object.image 61 | # image = tf.image.adjust_gamma(tf.cast(image_object.image, tf.float32), gamma=1, gain=1) # Scale image to (0, 1) 62 | label = image_object.label 63 | filename = image_object.filename 64 | 65 | if(if_random): 66 | min_fraction_of_examples_in_queue = 0.4 67 | min_queue_examples = int(TRAINING_SET_SIZE * min_fraction_of_examples_in_queue) 68 | print("Filling queue with %d images before starting to train. " "This will take a few minutes." % min_queue_examples) 69 | num_preprocess_threads = 1 70 | image_batch, label_batch, filename_batch = tf.train.shuffle_batch( 71 | [image, label, filename], 72 | batch_size = BATCH_SIZE, 73 | num_threads = num_preprocess_threads, 74 | capacity = min_queue_examples + 3 * BATCH_SIZE, 75 | min_after_dequeue = min_queue_examples) 76 | return image_batch, label_batch, filename_batch 77 | else: 78 | image_batch, label_batch, filename_batch = tf.train.batch( 79 | [image, label, filename], 80 | batch_size = BATCH_SIZE, 81 | num_threads = 1) 82 | return image_batch, label_batch, filename_batch 83 | 84 | 85 | def weight_variable(shape): 86 | initial = tf.truncated_normal(shape, stddev=0.05) 87 | return tf.Variable(initial) 88 | 89 | def bias_variable(shape): 90 | initial = tf.constant(0.02, shape=shape) 91 | return tf.Variable(initial) 92 | 93 | def conv2d(x, W): 94 | return tf.nn.conv2d(x, W, strides=[1, 1, 1, 1], padding='SAME') 95 | 96 | def max_pool_2x2(x): 97 | return tf.nn.max_pool(x, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='SAME') 98 | 99 | ## constructing the network. 100 | def malware_inference(image_batch): 101 | W_conv1 = weight_variable([5, 5, 3, 32]) 102 | b_conv1 = bias_variable([32]) 103 | 104 | x_image = tf.reshape(image_batch, [-1, IMAGE_SIZE, IMAGE_SIZE, 3]) 105 | 106 | h_conv1 = tf.nn.relu(conv2d(x_image, W_conv1) + b_conv1) 107 | h_pool1 = max_pool_2x2(h_conv1) # 112 108 | 109 | W_conv2 = weight_variable([5, 5, 32, 64]) 110 | b_conv2 = bias_variable([64]) 111 | 112 | h_conv2 = tf.nn.relu(conv2d(h_pool1, W_conv2) + b_conv2) 113 | h_pool2 = max_pool_2x2(h_conv2) # 56 114 | 115 | W_conv3 = weight_variable([5, 5, 64, 128]) 116 | b_conv3 = bias_variable([128]) 117 | 118 | h_conv3 = tf.nn.relu(conv2d(h_pool2, W_conv3) + b_conv3) 119 | h_pool3 = max_pool_2x2(h_conv3) # 28 120 | 121 | W_conv4 = weight_variable([5, 5, 128, 256]) 122 | b_conv4 = bias_variable([256]) 123 | 124 | h_conv4 = tf.nn.relu(conv2d(h_pool3, W_conv4) + b_conv4) 125 | h_pool4 = max_pool_2x2(h_conv4) # 14 126 | 127 | W_conv5 = weight_variable([5, 5, 256, 256]) 128 | b_conv5 = bias_variable([256]) 129 | 130 | h_conv5 = tf.nn.relu(conv2d(h_pool4, W_conv5) + b_conv5) 131 | h_pool5 = max_pool_2x2(h_conv5) # 7 132 | 133 | # with tf.variable_scope("fc"): 134 | W_fc1 = weight_variable([7*7*256, 2048]) 135 | b_fc1 = bias_variable([2048]) 136 | 137 | h_pool5_flat = tf.reshape(h_pool5, [-1, 7*7*256]) 138 | h_fc1 = tf.nn.relu(tf.matmul(h_pool5_flat, W_fc1) + b_fc1) 139 | 140 | h_fc1_drop = tf.nn.dropout(h_fc1, 1.0) 141 | 142 | W_fc2 = weight_variable([2048, 256]) 143 | b_fc2 = bias_variable([256]) 144 | 145 | h_fc2 = tf.nn.relu(tf.matmul(h_fc1_drop, W_fc2) + b_fc2) 146 | 147 | W_fc3 = weight_variable([256, 64]) 148 | b_fc3 = bias_variable([64]) 149 | 150 | h_fc3 = tf.nn.relu(tf.matmul(h_fc2, W_fc3) + b_fc3) 151 | 152 | W_fc4 = weight_variable([64, 5]) 153 | b_fc4 = bias_variable([5]) 154 | 155 | y_conv = tf.nn.softmax(tf.matmul(h_fc3, W_fc4) + b_fc4) 156 | # y_conv = tf.matmul(h_fc3, W_fc4) + b_fc4 157 | 158 | return y_conv 159 | 160 | 161 | def malware_train(): 162 | image_batch_out, label_batch_out, filename_batch = malware_input(if_random = False, if_training = True) 163 | 164 | image_batch_placeholder = tf.placeholder(tf.float32, shape=[BATCH_SIZE, 224, 224, 3]) 165 | image_batch = tf.reshape(image_batch_out, (BATCH_SIZE, IMAGE_SIZE, IMAGE_SIZE, 3)) 166 | 167 | label_batch_placeholder = tf.placeholder(tf.float32, shape=[BATCH_SIZE, 5]) 168 | label_offset = -tf.ones([BATCH_SIZE], dtype=tf.int64, name="label_batch_offset") 169 | label_batch_one_hot = tf.one_hot(tf.add(label_batch_out, label_offset), depth=5, on_value=1.0, off_value=0.0) 170 | 171 | logits_out = malware_inference(image_batch_placeholder) 172 | loss = tf.reduce_sum(tf.nn.softmax_cross_entropy_with_logits(labels=label_batch_one_hot, logits=logits_out)) 173 | # loss = tf.losses.mean_squared_error(labels=label_batch_placeholder, predictions=logits_out) 174 | 175 | train_step = tf.train.GradientDescentOptimizer(0.007).minimize(loss) ## 0.007 is the learning rate. 176 | 177 | saver = tf.train.Saver() 178 | 179 | with tf.Session() as sess: 180 | # Visualize the graph through tensorboard. 181 | file_writer = tf.summary.FileWriter("./logs", sess.graph) 182 | 183 | sess.run(tf.global_variables_initializer()) 184 | # saver.restore(sess, "C:/Users/admin/Project IAS/checkpoint-train.ckpt") 185 | coord = tf.train.Coordinator() ## A coordinator for threads implements a mechanism to coordinate the termination of a set of threads. 186 | threads = tf.train.start_queue_runners(coord=coord, sess = sess) ## Start all the queue runners collected in the graph. 187 | 188 | for i in range(TRAINING_SET_SIZE ):#* 100): 189 | image_out, label_out, label_batch_one_hot_out, filename_out = sess.run([image_batch, label_batch_out, label_batch_one_hot, filename_batch]) 190 | 191 | _, infer_out, loss_out = sess.run([train_step, logits_out, loss], feed_dict={image_batch_placeholder: image_out, label_batch_placeholder: label_batch_one_hot_out}) 192 | 193 | print(i) 194 | print(image_out.shape) 195 | print("label_out: ") 196 | print(filename_out) 197 | print(label_out) 198 | print(label_batch_one_hot_out) 199 | print("infer_out: ") 200 | print(infer_out) 201 | print("loss: ") 202 | print(loss_out) 203 | if(i%100 == 0): ## change to 50 204 | saver.save(sess, "C:/Users/admin/Project IAS/checkpoint-train.ckpt") 205 | 206 | coord.request_stop() 207 | coord.join(threads) 208 | sess.close() 209 | 210 | 211 | 212 | def malware_eval(): 213 | image_batch_out, label_batch_out, filename_batch = malware_input(if_random = False, if_training = False) 214 | 215 | image_batch_placeholder = tf.placeholder(tf.float32, shape=[BATCH_SIZE, 224, 224, 3]) 216 | image_batch = tf.reshape(image_batch_out, (BATCH_SIZE, IMAGE_SIZE, IMAGE_SIZE, 3)) 217 | 218 | label_tensor_placeholder = tf.placeholder(tf.int64, shape=[BATCH_SIZE]) 219 | label_offset = -tf.ones([BATCH_SIZE], dtype=tf.int64, name="label_batch_offset") ## tf.ones create a tensor mentioned shape(1st arg) and type(2nd arg) of 1's. 220 | label_batch = tf.add(label_batch_out, label_offset) 221 | 222 | logits_out = tf.reshape(malware_inference(image_batch_placeholder), [BATCH_SIZE, 5]) 223 | logits_batch = tf.to_int64(tf.arg_max(logits_out, dimension = 1)) 224 | 225 | correct_prediction = tf.equal(logits_batch, label_tensor_placeholder) 226 | accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32)) 227 | 228 | saver = tf.train.Saver() 229 | 230 | with tf.Session() as sess: 231 | sess.run(tf.global_variables_initializer()) 232 | saver.restore(sess, "C:/Users/admin/Project IAS/checkpoint-train.ckpt") 233 | coord = tf.train.Coordinator() ## A coordinator for threads implements a mechanism to coordinate the termination of a set of threads. 234 | threads = tf.train.start_queue_runners(coord=coord, sess = sess) 235 | 236 | accuracy_accu = 0.0 237 | 238 | for i in range(30): 239 | image_out, label_out, filename_out = sess.run([image_batch, label_batch, filename_batch]) 240 | 241 | accuracy_out, logits_batch_out = sess.run([accuracy, logits_batch], feed_dict={image_batch_placeholder: image_out, label_tensor_placeholder: label_out}) 242 | accuracy_accu += accuracy_out 243 | 244 | print(i) 245 | # print(image_out.shape) 246 | # print("label_out: ") 247 | # print(filename_out) 248 | # print(label_out) 249 | # print(logits_batch_out) 250 | # print (accuracy_accu, accuracy_out) 251 | 252 | print("Accuracy: ") 253 | print(accuracy_accu/30) 254 | 255 | coord.request_stop() 256 | coord.join(threads) 257 | sess.close() 258 | 259 | # malware_train() 260 | malware_eval() -------------------------------------------------------------------------------- /ias_full.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from PIL import Image 3 | import sys 4 | import glob 5 | import errno 6 | import os 7 | result=[] 8 | temp1=np.zeros((256,256)) 9 | image = np.zeros((256,256)) 10 | i,j,p=0,0,0 11 | path = "C:/Users/admin/Project IAS/Module1/data1/" ## Absolute path for the database. 12 | dirs = os.listdir( path ) ## gets all the directories in the given folder. 13 | path2 = "C:/Users/admin/Project IAS/Module1/image256/" ## Absolute path for storing 256x256 image. 14 | path3 = "C:/Users/admin/Project IAS/Module1/image36/" ## Absolute path for storing 36x36 image. 15 | for fil in dirs: 16 | image = np.zeros((256,256)) 17 | with open(path+fil, "rb") as f: ## read byte mode 18 | byte = f.read(1) ## reading 1 byte at a time. returns a byte object. 19 | i,j,p=0,0,0 20 | print(fil) 21 | #print (byte) 22 | while byte != b"": ## itereating till the end of the file 23 | # print (ord(byte),byte) ## ord gives ASCII value of that char. 24 | if(i<256): 25 | if(byte.decode("utf-8")==' ' or byte.decode("utf-8")=='\r' or byte.decode("utf-8")=='\n' or byte.decode("utf-8")=='?' ): 26 | byte = f.read(1) 27 | continue 28 | # image[i][p]=0 29 | #print (byte.decode("utf-8"),int(byte.decode("utf-8"),16),i) 30 | if(j>8): 31 | image[i][p]= int(byte.decode("utf-8"),16) ## byte.decode gives the value of the byte object which is converted into hexa-decimal. 32 | p=p+1 33 | j=j+1 34 | if(p>255): 35 | p=0 36 | j=0 37 | i=i+1 38 | byte = f.read(1) 39 | a=np.matrix(image) 40 | # a=np.matrix(image) 41 | print (a) 42 | result.append(np.array_equal(temp1,a)) 43 | temp1=a 44 | img = Image.fromarray(image, 'L') ## makes 2D array into an image. L-> The kind of pixel (8-bit pixel black and white) 45 | #img.show() 46 | img.save(path2+fil[0:-6]+".bmp") 47 | img = img.resize((36,36),Image.ANTIALIAS) 48 | img.save(path3+fil[0:-6]+".bmp") 49 | #img.show() 50 | print (result) 51 | 52 | 53 | 54 | # with open('C:/Users/admin/Project IAS/data/sample/0fHVZKeTE6iRb1PIQ4au.bytes', 'r') as file1: 55 | # print(file1.read()) 56 | # import pandas as pd 57 | # df = pd.read_csv("C:/Users/admin/Project IAS/data/sample/0fHVZKeTE6iRb1PIQ4au.bytes") 58 | # print(df.drop([0])) 59 | # with open('C:/Users/admin/Project IAS/data/sample/0gDsIvrylX5fPbG7cSBn.bytes', 'r') as file2: 60 | # print(file2.read()) 61 | ''' 62 | #same.discard('\n') 63 | 64 | # with open('some_output_file.txt', 'w') as file_out: 65 | # for line in same: 66 | # file_out.write(line) 67 | 68 | with open("apt1/APT1/VirusShare_00dbb9e1c09dbdafb360f3163ba5a3de.json", "rb") as f: 69 | byte = f.read(1) 70 | while byte: 71 | #print ord(byte) 72 | if(i<256): 73 | image[i][j]=ord(byte) 74 | j=j+1 75 | if(j>255): 76 | j=0 77 | i=i+1 78 | byte = f.read(1) 79 | 80 | print(np.matrix(image)) 81 | img = Image.fromarray(image, 'L') 82 | img.show() 83 | img = img.resize((36,36),Image.ANTIALIAS) 84 | img.show() 85 | 86 | ''' 87 | --------------------------------------------------------------------------------