import idaapi
import os
import re
import shutil
import idc


"""
PYSYM.PY
PURPOSE:TO SIMPLIFY THE PDB LOADING PROCESS

IDEA:	SIMPLE SCRIPT TO DLOAD THE SYMBOL (FROM MS), AND AUTO-LOAD
	IT INTO IDA.  NOTHING CRAZY, BUT A QUICK TIMESAVE TOOL TO HAVE.

USAGE:	INSTALL MICROSOFT DEBUGGING TOOLS (OR, JUST SYMCHK.EXE).
	LOAD THAT PATH INTO THE symchk VARIABLE.
	INSTALL IDAPython.
	RUN pysym.py

NOTE: 	In IDAPython .8, 'RunPlugin()' does not work.  So, once pysym.py
	is finished, just go back to 'file -> load file -> .pdb' and it 
	will already be in the working directory.

NOTE-2:	In order to use symchk, you must have the actual file being
	disasm'd.  Pysym.py assumes it's in the same directory as the idb.
	If this is not the case, make a copy of the file and place into the
	same dir as the idb.
"""

symchk = "C:\Program Files\Debugging Tools for Windows\symchk.exe"
load_pdb = 1	#SET TO 1 TO LOAD THE PDB ONCE GRABBED
load_dbg = 1	#SET TO 1 TO LOAD THE DBG ONCE GRABBED


def check_deps(input_file):
	if(os.path.isfile(input_file)):
		if(os.path.isfile(symchk)):
			return 0
		else:
			print 'ERROR: SYMCHK.EXE WAS NOT FOUND'
			return 1
	else:
		print 'ERROR: THE ORIGINAL FILE IS NOT IN THE SAME AS WHEN FIRST DISASSEMBLED'
		return 1

def getSymbols(file):
	path = os.path.split(file)[0]
	file = os.path.split(file)[1]
	os.chdir(path)
	arg = ' /r ' + file +  ' /s SRV*.\*http://msdl.microsoft.com/download/symbols '
	os.system('"' + symchk + '"' + arg)
	
def clean_folder():
	path = os.path.abspath('.')
	dirs = os.listdir('.')
	i = 0
	while(i < len(dirs)):
		if(os.path.isdir(dirs[i])):
			if(re.findall("\.((pdb)|(dbg))$",dirs[i]) != []):
				os.chdir(dirs[i])
				dirs2 = os.listdir('.')
				j = 0
				while(j < len(dirs2)):
					if(os.path.isdir(dirs2[j])):
						os.chdir(dirs2[j])
						dirs3 = os.listdir('.')
						k = 0
						while(k < len(dirs3)):
							if(os.path.isfile(dirs3[k])):
								if(re.findall("\.((pdb)|(dbg))$",dirs3[k]) != []):
									shutil.copyfile(dirs3[k],'../../'+dirs3[k]+'1')
							k = k + 1
					j = j + 1
				os.chdir(path)
				shutil.rmtree(dirs[i],ignore_errors=True,onerror=None)
				if(os.path.isfile('.\\pingme.txt')):
					os.remove('.\\pingme.txt')
		i = i + 1
		temp_files = os.listdir('.')
		f = 0
		while(f < len(temp_files)):
			if(re.findall("\.((pdb)|(dbg))1$",temp_files[f]) != []):
				os.rename(temp_files[f],temp_files[f][0:len(temp_files[f-1])])
			f = f + 1


print '\n\n\n         PySym\n  Andre Derek Protas\n--------------------------'
input_file =  GetInputFilePath()
if(check_deps(input_file) == 0):
	try:
		file_name = os.path.split(input_file)[1]
		print '~pysym: grabbing symbols'
		getSymbols(input_file)
		print '~pysym: joining files'
		clean_folder()
		print '~pysym: running file loaders'
		if(load_pdb == 1):
			pdb_file = re.sub("\..+$",".pdb",file_name)
			if(os.path.isfile(pdb_file)):
				print 'READY TO LOAD SYMBOLS'
				print 'File -> Load File -> PDB'
				#ONCE SUPPORTED, UNCOMMENT THESE TWO LINES
				#print '~pysym: running pdb loader'
				#idc.RunPlugin("pdb",0)
			else:
				print '!pysym: sorry, no pdb found'
		if(load_dbg == 1):
			pdb_file = re.sub("\..+$",".dbg",file_name)
			if(os.path.isfile(pdb_file)):
				print 'READY TO LOAD DBG'
				print 'File -> Load File -> DBG'
				#ONCE SUPPORTED, UNCOMMENT THESE TWO LINES
				#print '~pysym: running dbg loader'
				#idc.RunPlugin("mc12dbg.plw",0)
			else:
				print '!pysym: sorry, no dbg found'
		print '~pysym: c-ya'
	except Exception,inst:
		print '!pysym error: ',type(inst)
		print '~pysym: sorry'