main.py

import os
import argparse
import yara
from scanner.scan import MalwareExtractor
from scanner.extract_info import FileExtract
from utils.helperFun import write_results_to_file


if __name__ == 'main':
	parser = argparse.ArgumentParser()
	parser.add_argument('-rd', '--root_directory', help='Root directory to be scanned.')
	parser.add_argument('-yr', '--yara_rules', help='Compiled Yara rules file. Provide full path. Compiled file will be stored in rules directory.')
	parser.add_argument('-yd', '--yara_directory', help='Directory with Yara rules to be compiled. Output will be stored in output directory.')
	parser.add_argument('-fs', '--file_to_scan', help='Single file to be scanned and extract static data. Provide full path.')
	parser.add_argument('-fe', '--file_to_extract', help='Single file to extract static data. Provide full path.')

	args = parser.parse_args()

	if args.yara_directory:
		assert os.path.isdir(args.yara_directory)
		if not os.path.isdir('rules'):
			os.mkdir('rules')
		MalwareExtractor.compile_yara_rules(args.yara_directory)
		print("Yara rules compilation completed.")

	elif args.file_to_scan and args.yara_rules:
		try:
			assert os.path.isfile(args.file_to_scan)
			assert os.path.isfile(args.yara_rules)
			rules = yara.load(args.yara_rules)
			match = MalwareExtractor.scan_file(args.file_to_scan, rules)
			print("File scan completed.")
			mal_data = MalwareExtractor.extract_info_from_file(match)
			write_results_to_file(mal_data['data'], mal_data['file'], mal_data['match'])
		except:
			print("Please check the files inputted.")

	elif args.root_directory and args.yara_rules:
		if os.path.isdir(args.root_directory) and os.path.isfile(args.yara_rules):
			extractor = MalwareExtractor(args.root_directory, args.yara_rules)
			matches = extractor.scan_directory_struct()
			mal_data = MalwareExtractor.extract_info_from_files(matches)
			for mal in mal_data:
				write_results_to_file(mal['data'], mal['file'], mal['match'])

		else:
			print("Check that the directory and the Yara file exist.")
	elif args.file_to_extract:
		mal_data = FileExtract.get_info(args.file_to_extract)
		write_results_to_file(mal_data, args.file_to_extract)

	else:
		print("Please provide valid arguments. Run python main.py -h to see commands available")