exe_to_img.py

import argparse
import png
import os

# Creating arguments that can be called from the command line
# Allows for same code to be used for multiple conversion segregations with same script file
parser = argparse.ArgumentParser(prog="exe-to-img", description = "EXE to PNG conversion program")
parser.add_argument('-f', '--folder', type = str, help = "Source folder to operate on", default = "./")
parser.add_argument('-t', '--target', type = str, help = "Target folder to store generated images", default = "./")
parser.add_argument('-c', '--count', type = int, help = "Start value for image counter", default = 0)
fileType = parser.add_mutually_exclusive_group()
fileType.add_argument('-b', '--benign', action = 'store_true', help = "Mark coverted images as benign")
fileType.add_argument('-m', '--malware', action = 'store_true', help = "Mark coverted images as malware")

args = parser.parse_args()

targetFile = ""
if args.benign:
	targetFile += "ben"
elif args.malware:
	targetFile += "mal"

scanFolder = args.folder

# Counting total executables for realtime user feedback
print("Counting files, please wait...")
fileCount = 0
for root, dirs, files in os.walk(scanFolder):
	for file in files:
		if len(file.split('.')) > 1:
			ext = file.split('.')[1]
		else:
			ext = ""
		if ext.lower() != 'exe':
			continue
		fileCount += 1

# Looping through all files in the source folder
print("Files counted!", fileCount, "found")
print("Converting files to images and storing in [" + args.target + "] folder")
count = args.count
for root, dirs, files in os.walk(scanFolder):
	for file in files:
		if len(file.split('.')) > 1:
			ext = file.split('.')[1]
		else:
			ext = ""
		if ext.lower() != 'exe':
			continue
		
		# Reading and storing binary data from executable file
		print("Converting file: " + file + " (" + str(count + 1 - args.count) + " of " + str(fileCount) + ")")
		with open(root + "/" + file, "rb") as f:
			exeData = f.read()
		l = len(exeData)

		# Initializing matrix containing pixel data for output file
		# Target image dimensions will be 512 x 512 pixels to optimise data stored per sample and training speed
		# Matrix dimensions will be 512 x (512 * 3 bytes per pixel)
		pixMat = []
		for i in range(512):
			matRow = []
			for j in range(1536):
				if i * 512 + j >= l:
					matRow.append(0)						# Create black pixels in case of overflow
				else:
					matRow.append(exeData[i * 512 + j])		# Copy byte data to matrix for RGB values
			pixMat.append(matRow)
			matRow = []
		
		# Writing Pixel data matrix to image
		with open(args.target + targetFile + "_" + str(count).zfill(5) + ".png", "wb") as outFile:
			w = png.Writer(512, 512, greyscale = False, bitdepth = 8)
			w.write(outFile, pixMat)

		count = count + 1

exit(0)