checkpe.py

import shutil
from pprint import pprint
import pefile
import os
import array
import math
import pickle
import joblib
import sys
import argparse
from glob import glob
from overrides import sum3, len3
from typing import Dict
from collections import OrderedDict
import win32con, win32api
from var_dump import var_dump
import yara
import magic
import vt

print(f'{yara.__version__}')

pe: pefile.PE
suspect_log = 'suspect/list.log'
res = {}
apikey = 'f2a6dc6bea6de5bddb72f45313423e8ade04ac4a4da68afdbdc97d0c6703e885'

def get_entropy(data):
	if len(data) == 0:
		return 0.0
	occurences = array.array('L', [0] * 256)
	for x in data:
		occurences[x if isinstance(x, int) else ord(x)] += 1
	
	entropy = 0
	for x in occurences:
		if x:
			p_x = float(x) / len(data)
			entropy -= p_x * math.log(p_x, 2)
	
	return entropy

def get_resources(pe):
	"""Extract resources :
	[entropy, size]"""
	resources = []
	if hasattr(pe, 'DIRECTORY_ENTRY_RESOURCE'):
		try:
			for resource_type in pe.DIRECTORY_ENTRY_RESOURCE.entries:
				if hasattr(resource_type, 'directory'):
					for resource_id in resource_type.directory.entries:
						if hasattr(resource_id, 'directory'):
							for resource_lang in resource_id.directory.entries:
								data = pe.get_data(
									resource_lang.data.struct.OffsetToData,
									resource_lang.data.struct.Size)
								size = resource_lang.data.struct.Size
								entropy = get_entropy(data)
								resources.append([entropy, size])
		except Exception as e:
			return resources
	return resources

def get_version_info(pe):
	"""Return version infos"""
	res = {}
	for fileinfo in pe.FileInfo:
		if fileinfo.Key == 'StringFileInfo':
			for st in fileinfo.StringTable:
				for entry in st.entries.items():
					res[entry[0]] = entry[1]
		elif fileinfo.Key == 'VarFileInfo':
			for var in fileinfo.Var:
				res[var.entry.items()[0][0]] = var.entry.items()[0][1]
	if hasattr(pe, 'VS_FIXEDFILEINFO'):
		fill_fixed_info(pe, res)
	return res

def fill_fixed_info(pe, res):
	res['flags'] = pe.VS_FIXEDFILEINFO.FileFlags
	res['os'] = pe.VS_FIXEDFILEINFO.FileOS
	res['type'] = pe.VS_FIXEDFILEINFO.FileType
	res['file_version'] = pe.VS_FIXEDFILEINFO.FileVersionLS
	res['product_version'] = pe.VS_FIXEDFILEINFO.ProductVersionLS
	res['signature'] = pe.VS_FIXEDFILEINFO.Signature
	res['struct_version'] = pe.VS_FIXEDFILEINFO.StrucVersion

def extract_infos(fpath):
	global pe, res
	
	if res != None:
		del res
	res = {}
	
	try:
		pe = pefile.PE(fpath)
	except Exception as e:
		print(f'{e}', end='')
		return -1
	
	res['Machine'] = pe.FILE_HEADER.Machine
	res['SizeOfOptionalHeader'] = pe.FILE_HEADER.SizeOfOptionalHeader
	res['Characteristics'] = pe.FILE_HEADER.Characteristics
	res['MajorLinkerVersion'] = pe.OPTIONAL_HEADER.MajorLinkerVersion
	res['MinorLinkerVersion'] = pe.OPTIONAL_HEADER.MinorLinkerVersion
	res['SizeOfCode'] = pe.OPTIONAL_HEADER.SizeOfCode
	res['SizeOfInitializedData'] = pe.OPTIONAL_HEADER.SizeOfInitializedData
	res['SizeOfUninitializedData'] = pe.OPTIONAL_HEADER.SizeOfUninitializedData
	res['AddressOfEntryPoint'] = pe.OPTIONAL_HEADER.AddressOfEntryPoint
	res['BaseOfCode'] = pe.OPTIONAL_HEADER.BaseOfCode
	try:
		res['BaseOfData'] = pe.OPTIONAL_HEADER.BaseOfData
	except AttributeError:
		res['BaseOfData'] = 0
	res['ImageBase'] = pe.OPTIONAL_HEADER.ImageBase
	res['SectionAlignment'] = pe.OPTIONAL_HEADER.SectionAlignment
	res['FileAlignment'] = pe.OPTIONAL_HEADER.FileAlignment
	res['MajorOperatingSystemVersion'] = pe.OPTIONAL_HEADER.MajorOperatingSystemVersion
	res['MinorOperatingSystemVersion'] = pe.OPTIONAL_HEADER.MinorOperatingSystemVersion
	res['MajorImageVersion'] = pe.OPTIONAL_HEADER.MajorImageVersion
	res['MinorImageVersion'] = pe.OPTIONAL_HEADER.MinorImageVersion
	res['MajorSubsystemVersion'] = pe.OPTIONAL_HEADER.MajorSubsystemVersion
	res['MinorSubsystemVersion'] = pe.OPTIONAL_HEADER.MinorSubsystemVersion
	res['SizeOfImage'] = pe.OPTIONAL_HEADER.SizeOfImage
	res['SizeOfHeaders'] = pe.OPTIONAL_HEADER.SizeOfHeaders
	res['CheckSum'] = pe.OPTIONAL_HEADER.CheckSum
	res['Subsystem'] = pe.OPTIONAL_HEADER.Subsystem
	res['DllCharacteristics'] = pe.OPTIONAL_HEADER.DllCharacteristics
	res['SizeOfStackReserve'] = pe.OPTIONAL_HEADER.SizeOfStackReserve
	res['SizeOfStackCommit'] = pe.OPTIONAL_HEADER.SizeOfStackCommit
	res['SizeOfHeapReserve'] = pe.OPTIONAL_HEADER.SizeOfHeapReserve
	res['SizeOfHeapCommit'] = pe.OPTIONAL_HEADER.SizeOfHeapCommit
	res['LoaderFlags'] = pe.OPTIONAL_HEADER.LoaderFlags
	res['NumberOfRvaAndSizes'] = pe.OPTIONAL_HEADER.NumberOfRvaAndSizes
	
	# Sections
	res['SectionsNb'] = len(pe.sections)
	entropy = list(map(lambda x: x.get_entropy(), pe.sections))
	res['SectionsMeanEntropy'] = sum(entropy) / float(len(entropy))
	res['SectionsMinEntropy'] = min(entropy)
	res['SectionsMaxEntropy'] = max(entropy)
	
	raw_sizes = list(map(lambda x: x.SizeOfRawData, pe.sections))
	res['SectionsMeanRawsize'] = sum(raw_sizes) / float(len(raw_sizes))
	res['SectionsMinRawsize'] = min(raw_sizes)
	res['SectionsMaxRawsize'] = max(raw_sizes)
	virtual_sizes = list(map(lambda x: x.Misc_VirtualSize, pe.sections))
	res['SectionsMeanVirtualsize'] = sum(virtual_sizes) / float(
		len(virtual_sizes))
	res['SectionsMinVirtualsize'] = min(virtual_sizes)
	res['SectionMaxVirtualsize'] = max(virtual_sizes)
	
	# Imports
	try:
		res['ImportsNbDLL'] = len(pe.DIRECTORY_ENTRY_IMPORT)
		imports = list(sum((x.imports for x in pe.DIRECTORY_ENTRY_IMPORT), []))
		res['ImportsNb'] = len(imports)
		res['ImportsNbOrdinal'] = len(
			list(filter(lambda x: x.name is None, imports)))
	except AttributeError:
		res['ImportsNbDLL'] = 0
		res['ImportsNb'] = 0
		res['ImportsNbOrdinal'] = 0
	
	# Exports
	try:
		res['ExportNb'] = len(pe.DIRECTORY_ENTRY_EXPORT.symbols)
	except AttributeError:
		# No export
		res['ExportNb'] = 0
	# Resources
	resources = get_resources(pe)
	res['ResourcesNb'] = len(resources)
	if len(resources):
		entropy = list(map(lambda x: x[0], resources))
		res['ResourcesMeanEntropy'] = sum(entropy) / float(len(entropy))
		res['ResourcesMinEntropy'] = min(entropy)
		res['ResourcesMaxEntropy'] = max(entropy)
		sizes = list(map(lambda x: x[1], resources))
		res['ResourcesMeanSize'] = sum(sizes) / float(len(sizes))
		res['ResourcesMinSize'] = min(sizes)
		res['ResourcesMaxSize'] = max(sizes)
	else:
		res['ResourcesNb'] = 0
		res['ResourcesMeanEntropy'] = 0
		res['ResourcesMinEntropy'] = 0
		res['ResourcesMaxEntropy'] = 0
		res['ResourcesMeanSize'] = 0
		res['ResourcesMinSize'] = 0
		res['ResourcesMaxSize'] = 0
	
	# Load configuration size
	try:
		res['LoadConfigurationSize'] = pe.DIRECTORY_ENTRY_LOAD_CONFIG.struct.Size
	except AttributeError:
		res['LoadConfigurationSize'] = 0
	
	# Version configuration size
	try:
		version_infos = get_version_info(pe)
		res['VersionInformationSize'] = len(version_infos.keys())
	except AttributeError:
		res['VersionInformationSize'] = 0
	return res

def add_to_suspect(path):
	with open(suspect_log, 'a') as f:
		f.write("\n" + ("=" * 40) + f'\nFile: {path}\n')
		for match in yara_founds:
			f.write(f'yara: {match}\n')
		for key, value in res.items():
			f.write(f'pe->{key}: {value}\n')
	try:
		shutil.copy(path, 'suspect')
	except Exception as e:
		print(f'copy error {e}')

def is_pe(path):
	try:
		with open(path, 'rb') as f:
			return f.read(2) == b"MZ"
	except Exception as e:
		return False

def yara_setup(all=False):
	global yara_rules
	# yara.set_config(stack_size=131072)
	# yara.set_config(max_strings_per_rule=50000)
	# yara.set_config(max_match_data=2048)
	
	if all:
		yara_rules = yara.compile(r'f:\yara-malware-rules\index.yar')
	else:
		yara_rules = yara.compile(filepaths={'vm' : r'f:\yara-malware-rules\antidebug_antivm_index.yar',
		                                     'mw' : r'f:\yara-malware-rules\malware_index.yar',
		                                     'cve': r'f:\yara-malware-rules\cve_rules_index.yar',
		                                     'xpl': r'f:\yara-malware-rules\exploit_kits_index.yar'})

def erase_suspect():
	print(f'erasing suspect ...')
	
	try:
		files = glob('suspect/**')
		
		for file in files:
			try:
				win32api.SetFileAttributes(file, win32con.FILE_ATTRIBUTE_NORMAL)
				shutil.rmtree("suspect", ignore_errors=True)
			except Exception as e:
				print(f'setf({file}): {e}')
		
		if not os.path.isdir('suspect'):
			os.mkdir('suspect')
	except Exception as e:
		print(f'erase: {e}')

def process_file(file):
	print(f'scanning {file}')
	yara_setup(all=True)
	
	try:
		print(f'type: ', end='')
		m = magic.from_file(file)
		print(f'{m}')
	except Exception as e:
		print(f'magic exception: {e}')
	
	try:
		matches = yara_rules.match(file)
		for match in matches:
			if verbose:
				for m in match.strings:
					offset, _, string = m
					print(f'yara: {match} 0x{offset:08x} {string.decode("utf-8", errors="ignore")}')
			else:
				print(f'yara: {match}')
	except Exception as e:
		print(f'yara: {e}')
	
	# try:
	#     print(f'checking with VirusTotal ... ', end='')
	#     with vt.Client(apikey) as client:
	#         with open(file, encoding='utf-8', errors='ignore') as f:
	#             analysis = client.scan_file(file=f, wait_for_completion=True)
	#             print(f'file uploaded.')
	#             pprint(analysis)
	# except Exception as e:
	#     print(f'vt: {e}')
	sys.exit(0)

if __name__ == '__main__':
	global yara_founds, yara_rules
	
	parser = argparse.ArgumentParser(description='Detect malicious files')
	parser.add_argument('-v', action='store_true', help='verbose', default=False, required=False)
	parser.add_argument('-d', help='File dir to be tested')
	parser.add_argument('-e', action='store_true', required=False, default=False,
	                    help='erase suspect dir prior to scanning')
	parser.add_argument('-f', required=False, default=False, help="scan file")
	args = parser.parse_args()
	
	yara_setup()
	
	verbose = args.v
	
	if args.f:
		process_file(args.f)
	
	if args.e:
		erase_suspect()
	
	# Load classifier
	p = os.path.dirname(os.path.realpath(__file__))
	
	print(f'classifier: {p}')
	clf = joblib.load(os.path.join(p, r'classifier\classifier.pkl'))
	
	if os.path.exists(suspect_log):
		os.unlink(suspect_log)
	
	with open(os.path.join(p, r'classifier\features.pkl'), 'rb') as f:
		features = pickle.load(f)
	
	print(f'loaded {len(features)} features')
	
	root = os.path.join(str(args.d), "**.dll")
	print(f'calculating {root} ...', end='')
	files = glob(root, recursive=True)
	print(f'{len(files)} files')
	
	for file in files:
		if os.path.isdir(file):
			continue
		
		print(f'{file} ... ', end='')
		
		yara_founds = []
		try:
			m = magic.from_file(file)
			print(f'{m}')
		except Exception as e:
			print(f'magic exception: {e}')
		
		if not is_pe(file):
			print('\tskip')
			continue
		
		try:
			matches = yara_rules.match(file)
			if len(matches):
				yara_founds = matches
				for match in yara_founds:
					print(f'[{match}] ', end='')
		except Exception as e:
			print(f'yara: {e}')
		
		res = extract_infos(file)
		if isinstance(res, dict):
			pe_features = list(map(lambda x: res[x], features))
			result = clf.predict([pe_features])[0]
			suspect = result == 1
			
			rep = ['OK', 'SUSPECT'][result]
			print(f'nn: {rep} ', end='')
			
			if suspect and len(yara_founds):
				print(f'\tflagged')
				add_to_suspect(file)
		else:
			print(f'\tskipped (no info)')
			continue