############################################################################
#
#	binStar: Python library to facilitate binary analysis
#
#
############################################################################

import json, os, sys, getopt, cPickle
from subprocess import call,Popen,PIPE,check_output


class binStar:
#---------------------------------Strings----------------------------------------------------------------------------------------------------------------------------
	def regexStringSearch(self,strings,word):
		"""
		Takes Perl-Regex and string list, searches for a match, returns results

		@param	list:		List of strings to search through
		@type	list:		list<str>
		@param 	pattern:	Perl-style regex to search for
		@type 	pattern: 	string
		@rtype:				list<long>
	    @return: 			List of addresses(longs) of strings matching regex
	    """
	   	matches = []
		for string in strings:
			if word in string:
				matches.append(string)

		return matches

	#This uses IDA functions, need to fix, replace "checkBuildStrings"
	def strings(self,target_file):
		"""
		Function to return list of all strings in target binary

		@rtype:		list<str>
	    @return:	List of strings within target binary
	    """
		strings_exe = os.path.expandvars(self.__cfg_dict__["cfg"]["strings_exe"])
		args = strings_exe + " -n 4 " + target_file
		out = check_output(args)
		strings = out.split('\n')
		return strings
		

	def stringSearch(self,target_file,strings_file):
		"""
		Searches binary for strings matching regex

		@param	regex:	Perl-style regex
		@type	regex:	str
		@rtype:			list<(str,long)>
		@return:		list of string,address tuples
		"""
		strings = self.strings(target_file)
		words_fh = open(strings_file,'r')
		words = words_fh.read()
		words = words.split('\n')
		target = os.path.basename(target_file)
		str_report = open(self.__work_dir__+target+"stringsReport.txt","w")
		for word in words:
			matches = self.regexStringSearch(strings,word)
			if len(matches) < 1:
				continue
			str_report.write("########################################\n")
			str_report.write(word+"\n########################################\n\n\n\n")
			for match in matches:
				str_report.write(match+'\n')
			str_report.write("\n\n\n\n\n")
		str_report.close()

	def pdbSymbol(self,regex):
		"""
		Searches PDB file for symbols matching regex

		@param	regex:	Perl-style regex
		@type	regex:	str
		@rtype:			list<(str,long)>
		@return:		list of string,address tuples
		"""
		symList = pdbSymbolList()
		addrs = regexStringSearch(symList,regex)
		return addrs

#----------------------------------PDB------------------------------------------------------------------------------------------------------------------------

	def loadPdbFile(self,pdb_file,base=0x140000000):
		"""
		Loads pdb file into a lookup class object

		@param	base:	Specifies the base address of the binary. Default value matches that of IDAPython
		@type	base:	long
		@return:		void
		"""									
		pdb = [(pdb_file,base)]						#Lookup takes a (pdb,baseAddress) tuple
		print "Loading PDB file. This may take some time..."
		goodLooks = Lookup(pdb)						#Instantiate global Lookup object
		print "PDB file loaded"
		return goodLooks

	def picklePdb(self,pdb_file,pickle_file):
		if os.path.isfile(pickle_file):
			print "Pickle file already exists, lets save some time"
			return
		pdbLookup = self.loadPdbFile(pdb_file)
		output = open(pickle_file,'wb')
		cPickle.dump(pdbLookup,output)
		output.close()

	def pdbSymbolList(self):
		"""
		Extracts symbol list from Lookup class object.

		@rtype:		list<str>
		@return:	String list of the PDB symbol names
		"""
		if (__pdb_file__ == None or goodLooks == None):	#Has the pdb file been loaded yet?
			loadPdbFile()
		symDict = goodLooks.names						#Grab dictionary of Symbol names
		symbolList = []
		for s in symDict.values():
			symbolList.extend(s)						#Add each name to symbol list
		return symbolList

#---------------------------------Analyze Binaries------------------------------------------------------------------------------------------------------------------------

	def runBin(self):
		"""
		Does all the work. For list of target_binaries:
			Loads target config
			If specified, runs binary and dumps module
			Performs strings check
			Pickles PDB file
			Runs IDA analysis
		"""
		target_binaries = self.__cfg_dict__["target_binaries"]
		dumpExe = os.path.expandvars(self.__cfg_dict__["cfg"]["dumpIt"])

		for idx,target_config in enumerate(target_binaries):
			if target_config == "":
				break
			target = target_config["target"]
			binary = os.path.expandvars(target_config["filepath"])
			args = target_config["cmd_line_args"]
			run_bin = target_config["dump_module"].lower() == "true"
			run_tests = target_config["run_tests"].lower() == "true"
			strings_file = os.path.expandvars(target_config["strings_file"])
			pdb_file = os.path.expandvars(target_config["pdb_file"])

			pickle_file = self.__work_dir__+target+".pkl"

			if run_bin:																#If config says run executable, dump live memory
				if target == "":
					target = os.path.basename(binary)								#If no target specified, the live binary itself is dumped
				self.__target_file__ = self.__work_dir__+target+".bin"				#File path to dump
				try:
					call([dumpExe,binary,args,target,self.__work_dir__])			#Runs dumpIt
				except:
					print "dumpIt call failed, quiting"
					sys.exit(1)
			else:
				self.__target_file__ = self.__work_dir__+os.path.basename(binary)	#If not dumping, set executable as target
				os.rename(binary,self.__target_file__)								#Makes local copy of target, to be safe

			if not os.path.isfile(self.__target_file__):
				print "Target file does not exist!"									#Checks if target file exists
				sys.exit(0)

			if run_tests:
				if strings_file != "":												#If tests are on, checks if string file provided
					self.stringSearch(self.__target_file__,strings_file)			#Runs the string search

				if not pdb_file == "":
					self.picklePdb(pdb_file,pickle_file)							#Pickles the pdb file, if PDB file is provided
				else:
					print "No symbol file specified"
				self.runIDA(idx)													#Sets up the IDA tests
			else:
				print "Binary tests skipped"

	def runIDA(self,idx):
		"""
		Runs IDA in autonomous mode with idaStar.py

		@param	idx:	Index of target in target_binaries
		@type	idx:	int
		"""
		ida = os.path.expandvars(self.__cfg_dict__["cfg"]["IDAw"])
		idaStar = os.path.expandvars(self.__cfg_dict__["cfg"]["idaStar"])

		cfg = self.__cfg_file__ + " " + str(idx)
		ida_args = "-a -A -S\""+idaStar+" "+cfg+"\" "+self.__target_file__		#Builds IDA cmd args: -a "Skip autoanalysis"
		ida_cmd = ida+' '+ida_args												#-A "Autonomous mode" -S'script.py arg1' "Script file to run, plus args"
		print "Running IDA analysis"
		try:
			p = check_output(ida_cmd)											#Runs command, captures output
		except:
			print "IDA call failed, quiting"
			sys.exit(1)
		print "IDA tests have completed"

#---------------------------------Class Utils------------------------------------------------------------------------------------------------------------------------

	def cfgParse(self,cfg_file):
		"""
		Takes file path, and loads JSON data into config dictionary
		Currently, the following entries are populated:
			'binary':	File path to target binary
			'opts': 	Command-line options to run binary with
			'pdb': 		File path to target's associated PDB file
			'strings':	File path to file containing newline delim'd regexes for regString()
			'bytes':	File path to file containing newline delim'd hex string byte patterns for byteSig()
			'output':	Directory to store working files and results report file

		@param 	cfg_file:	File path to json formatted config file
		@type 	pattern:	string
		@rtype:				dict
	    @return: 			Dictionary populated with data in config file
	    """	
		fp = open(cfg_file)
		cfg_dict = json.loads(fp.read().strip())
		fp.close()
		print "Config file has been loaded"
		return cfg_dict


	def __init__(self,argv):
		"""
		binStar class initializer
		"""
		usage = "binStar.py -c configuration_file"
		self.__cfg_dict__ = None
		self.__logFile__ = None
		self.__pdb_file__ = None
		try:
			opts,args = getopt.getopt(argv,"hc:",["config="])		#Parse command line arguments
		except getopt.GetoptError:
			print usage
			sys.exit(2)
		for opt, arg in opts:
			if opt == '-h':
				print "\nPlease provide a valid conig file to binStar. See the README for more information\n"
				print usage
				sys.exit()
			elif opt in ("-c","--config"):
				self.__cfg_file__ = arg
				self.__cfg_dict__ = self.cfgParse(arg)
			else:
				print usage
		if len(sys.argv)<2:
			print usage

		if self.__cfg_dict__ == None:
			print "No config file supplied, quiting"
			sys.exit(0)

		self.__work_dir__ = os.path.expandvars(self.__cfg_dict__["cfg"]["work_dir"])
		if not os.path.exists(self.__work_dir__):
			os.makedirs(self.__work_dir__)


#---------------------------------Main------------------------------------------------------------------------------------------------------------------------

if __name__ == "__main__":
	bs = binStar(sys.argv[1:])

	sys.path.append(os.path.expandvars(bs.__cfg_dict__["cfg"]["pylibs"]))	#Allows for lib location to be specified in config file
	from pdbparse.symlookup import Lookup

	bs.runBin()