automate.py

import magic
import glob
import os
import uuid

SAMPLES_FOLDER = '../radare2-regressions'

files = glob.glob(SAMPLES_FOLDER + '/**/*', recursive=True)

# print(files)

OUT_DIR = './images/'

CMD = './dump2png -h 224 -w 224 -p color -o %s %s'

if not os.path.exists(OUT_DIR):
	os.makedirs(OUT_DIR)

counters = {'ELF':0, 'PE':0, 'DOS':0, 'Dalvik':0, 'MachO':0}

def make_image(p, m, label):
	u = '%s_%d.png' %(label, counters[label])
	filename = os.path.join(OUT_DIR, u)
	print(CMD%(filename, p))
	os.system(CMD%(filename, p))
	counters[label] = counters[label] + 1

for f in files:
	if os.path.isdir(f):
		continue
	m = None
	try:
		m = magic.from_file(f)
	except Exception as e:
		continue
	if 'ELF' in m:
		pass
		# print(m)
		make_image(f, m, 'ELF')
	elif 'PE' in m:
		make_image(f, m, 'PE')
	elif 'DOS' in m:
		make_image(f, m, 'DOS')
	elif 'Dalvik' in m:
		make_image(f, m, 'Dalvik')
	elif 'Mach' in m:
		make_image(f, m, 'MachO')
	else:
		print(m)