""" DisassSubRoutine is an object which holds a subroutine, as parsed by disass-ng.py
It expects to be called by a "disass" object, which keeps track of the entire binary
Calls are made to self.parent.SUBS, self.parent.SYMBOLS, etc....

"""
import sys
import sre
import os

SRCH_ASMLINE = '^\s*[0-9a-fA-F]{4,8}:'
# Sub Parser Constants
SRCH_SUBTITLE = "^[0-9a-fA-F]{7,8}\s*<.*>:"
SRCH_PUSHEBP = "55.*push\s+%ebp"
SRCH_MOVESPEBP = "89.*[eE]5.*mov\s+%esp,\s*%ebp"
SRCH_RET = "[cC]3.*ret"
SRCH_CALL = "call "
SRCH_JMP = "j[abezlgnm][ezps]{0,1} "
SRCH_MEMLOC = "0x[0-9a-fA-F][0-9a-fA-F]{4,7}[^(]"
CONST_ENDSUB = "\nRETURN\n"

def chopaddress(address):
	addy = address.strip(' \*$')
	try:
		if (addy[1] == 'x'):
			addy = addy[2:]
		if (addy[0] == '0'):
			addy = addy[1:]
	except:
		pass
	return addy
			


class DisassSubRoutine:

	def __init__(self, parent = None, address = None):
		self.parent = parent
		self.startaddress = address
		self.endaddress = ""
		self.name = ""
		self.lines = []
		self.value = []
		self.vrbls = []
		self.callers = {}
		self.jmpers = {}

		#print >>sys.stderr,("New Sub")
		
	def __str__(self):
		retvar = "\n\n"
		retvar += "\nSubroutine: %s \t %d lines" %( self.name,len(self.lines))
		
		# Print out Local Variables and Sub Parms
		if (len(self.vrbls) > 0):
			retvar += "\nVariables:"
			vrb = self.vrbls
			vrbls = []
			for i in vrb:
				vrbls.append(long(i.split('(')[0],16))
			vrbls.sort()
			for x in range(len(vrbls)):
				size = 0
				var1 = vrbls[x]
				if (var1 <= 2000000):
					if (x < len(vrbls)-1):
						var2 = vrbls[x+1]
					else:
						var2 = 0
					# The current variable is a Parameter
					if (var2 <= 2000000):
						size = var2 - var1
					else:
						size = 4
				else:
					# The current variable is a local var
					if (x < len(vrbls)-1):
						var2 = vrbls[x+1]
					else:
						var2 = 0x100000000
					size = var2 - var1
				
				retvar += "\n    %16x (%5x) _____________________________________________________" % (vrbls[x],size)
			
		# Possibly temporary:  list beginning and end addresses:
		if (len(self.lines) > 7):
			retvar += "\nStarting address: %s\t\tEnding address: %s\n"%(self.startaddress,self.endaddress)
		# Print out "Callers" and "JMPers"
		if (len(self.callers) > 0):  
			retvar += "\nCalled By:\n"
			for x in self.callers.keys():
				retvar += "\t%s->\t%s\n"%(x, self.callers.get(x))
		
		if (len(self.jmpers) > 0):
			retvar += "\nJMPed into from:\n"
			for x in self.jmpers.keys():
				retvar += "\t%s:\t%s\n"%(x, self.jmpers.get(x))
	
		
		# Print the Disassembly
		for EACH in range(len(self.lines)):
			retvar += "\n%-85s\t%s" % (self.lines[EACH],self.value[EACH])

		retvar += "\nEnd Subroutine: %s \t %d lines   " % (self.name,len(self.lines))
		return retvar
	
	def setName(self,name = None):
		### Set the name first
		if (name):
			self.name = name
		else:
			#  Find Name among "peer" objects
			self.name = self.parent.getSubName(self.startaddress)
			#print >>sys.stderr, ("Naming %s::%s" % (self.startaddress,self.name))
			
		### Now set the start and end addresses and register this sub with the LOOKUPS table
		#  Find first line of ASM
		ptr = 0
		while (sre.search(SRCH_ASMLINE, self.lines[ptr]) == None and ptr< (len(self.lines)-1)):
			ptr += 1
		address = self.lines[ptr][:self.lines[ptr].find(":")]
		#chop off leading spaces... 
		self.startaddress = address.strip()
		self.parent.registerLookup(address,self.name)
		
		#  Find last line of ASM  (set boundaries)
		ptr = len(self.lines)-1
		while (sre.search(SRCH_ASMLINE, self.lines[ptr]) == None and ptr> 0):
			ptr -= 1
		address = self.lines[ptr][:self.lines[ptr].find(":")]
		self.endaddress = address.strip()

		#print >>sys.stderr, ("NAMING SUB: %s"%self.name)
	
	def append(self, line):
		self.lines.append(line.strip())
		self.value.append("")
		#if (self.name == "" and len(self.lines) > 1):
			#print >>sys.stderr,("calling DisassSubRoutine.setName() with current sub length of %d\n %s\n  %s" % (len(self.lines), self.lines[0],self.lines[-1]))
			
		#self.setName()
	
	def process(self):
		""" This tells the object that all lines are in so it can run it's post-processing tasks"""
		print >>sys.stderr,("BEGIN Processing Sub: %s"%self.name)
		### Bring the lookups local ###
		self.LOOKUPS = self.parent.LOOKUPS
		self.DATA = self.parent.DATA
		self.SUBS = self.parent.SUBS
		self.LOOKUPKEYS = self.parent.LOOKUPKEYS			
			

		#self.setName()     #Already done
		print >>sys.stderr,("tagVars()")
		self.tagVars()
		print >>sys.stderr,("tagPLT")
		self.tagPLT()		
		print >>sys.stderr,("END Processing Sub: %s"%self.name)

	def getRepr(self,address,unknown="",findoffset=True):

		### Check self.LOOKUPS for the address ###
		addy = chopaddress(address)
		#print >>sys.stderr,(addy)
		name = self.LOOKUPS.get(addy)

		### Check self.LOOKUPS.get('DATA').getRepr()
		if (not name):
			if (self.DATA):
				name = self.DATA.getRepr(addy)

		if (name and len(name) > 0 or not findoffset):  return name

		### Do extended lookup  (previous address Repr  "+0x####")
                ### Next, Look through the Subroutines.  Is it IN one of them?  If so,
		s = 0
		sub = ""
		lclkeys = self.LOOKUPKEYS
                try:
                        # This will determine what local sub owns this address
                        addrint = int(addy,16)
                        # Check SubRoutines for location
                        subrange = ""
                        ##### NEEDS REWORKING OR DELETION.  Use counter, and compare between this and the next sub. #####
                        for s in range(len(lclkeys)-1):
				sub = lclkeys[s]
				if (sub < addy and lclkeys[s+1] > addy):
					subint = int(sub, 16)
					diff = addrint-subint
					#print >>sys.stderr,("disass.getRepr.subsub  :%x:%x:%s:"%(addrint,subint,sub))
					if (addrint >= subint):
						subrange = "(%s %+.4x)"%(self.LOOKUPS.get(sub),diff)
                        name += subrange
                        #print >>sys.stderr,("getRepr(): %s"%subrange)

                except Exception,e:
                        print >>sys.stderr,("getRepr(): %s : %d::%s"%(e,s,sub))
		if (not name):  name = unknown
		return name

				
	def addCaller(self, src, dest, name):
		# clean up passed in address, stripping it of any "0x" or " " 
		#sre.sub("(0x|\s)","",dest)
		dst = chopaddress(dest)
		#print >>sys.stderr,("SUB:  %s,    :%s:%s:%s:"%(self.name, self.startaddress, dest, self.endaddress))
		if (dst >= self.startaddress and dst <= self.endaddress):
			self.callers.setdefault(src,dst)
			return True
		return False

	def addJMPer(self, src, dest, name):
		# clean up passed in address, stripping it of any "0x" or " " 
		#sre.sub("(0x|\s)","",dest)
		dst = chopaddress(dest)
		#print >>sys.stderr,("SUB:  %s,    :%s:%s:%s:"%(self.name, self.startaddress, dest, self.endaddress))
		if (dst >= self.startaddress and dst <= self.endaddress):
			self.jmpers.setdefault(src,dst)
			return True
		return False

	def tagVars(self):
		for line in self.lines:
			for segment in line.split(','):
				varidx = -1
				varstart = -1
				varend = -1
				try:
					varidx = segment.find('(%ebp)')
					#print >>sys.stderr,("SubRoutine:tagVars:: varidx: %d"%varidx)

					if (varidx > -1):
						varstart = segment.index('0x')
						varend = segment.index('(%')+6
						var = segment[varstart:varend]
						try:
							self.vrbls.index(var)
						except:
							self.vrbls.append(var)

						#print >>sys.stderr,(" FoundVar:  %s:%s"%(line[varstart:varend+1],line))
				except Exception, e:
					print >>sys.stderr,("tagSubVars:  Something broke at %s" % (segment))
					print >>sys.stderr,(e)
					
	def tagPLT(self):
		#  precompile local regular expressions before looping.
		sre_memloc = sre.compile(SRCH_MEMLOC)
		sre_jmps = sre.compile(SRCH_JMP)

		####### Stage 1: Associate GOT and PLT entries and Tag disassassembly calls to PLT with the appropriate	name #######
		for EACH in range(len(self.lines)):
			thisline = self.lines[EACH]
			breakdown = []
			lineaddress = ""
			targetaddress = ""
			direction = ""
			cmd = ""
			
			callspot = thisline.find(SRCH_CALL)
			jmps = sre.findall(sre_jmps, thisline)
			memloc = sre.findall(sre_memloc, thisline)
			
			if (callspot > -1):
				#print >>sys.stderr, ("tagPLT::call")
				cmd = "call"
				breakdown.append(thisline[:callspot])
				breakdown.append(thisline[callspot+4:])
				
				lineaddress = breakdown[0][:breakdown[0].find(":")]
				lineaddress = lineaddress.strip()
				### This is a hack.. May have to track the memory location.
				### We'll try this for speed
				targetaddress = breakdown[1][3:]
				foundarrow = targetaddress.find("<")
				if (foundarrow >-1): targetaddress = targetaddress[:foundarrow].strip()

				self.parent.addBREAK("*0x%s" % lineaddress)
				
				try:
					name = "%s %s" % (cmd,self.getRepr(targetaddress))
				except:
					name = "call undef"
				
				self.value[EACH] += " %s (brkpt: %i)\n" % (name,len(self.parent.BREAK)-1)
				if (not (targetaddress >=self.startaddress and targetaddress <= self.endaddress)):
					self.parent.registerCall("%s:%s"%(self.name,lineaddress), targetaddress, name)			
				#print >>sys.stderr, ("tagPLT::call::END")
			elif (jmps):
				#print >>sys.stderr, ("tagPLT::jmp")
				jmpspot = thisline.find(jmps[0])
				cmd = thisline[jmpspot:]
				cmd = cmd.split()[0]

				breakdown.append(thisline[:jmpspot])
				breakdown.append(thisline[jmpspot+4:])
				
				lineaddress = breakdown[0][:breakdown[0].find(":")]
				lineaddress = lineaddress.strip()
				### This is a hack.. May have to track the memory location.
				### We'll try this for speed
				targetaddress = breakdown[1][3:]
				foundarrow = targetaddress.find("<")
				if (foundarrow>-1): targetaddress = targetaddress[:foundarrow].strip()
				
				# Register JMP with parent, and subsequently other subs
				if (not (targetaddress >=self.startaddress and targetaddress <= self.endaddress)):
					try:
						name = self.getRepr(targetaddress)
					except:
						name = "undef"
					self.parent.registerJMP("%s-%s"%(self.name,lineaddress), targetaddress, name)
				else:
					name = "local"
				try:
					targetaddress = chopaddress(targetaddress)
					lineaddress = chopaddress(lineaddress)
					direction = int(targetaddress,16) - int(lineaddress,16)
					dirstr = " %+.6X"%(direction)
				except:
					#print >>sys.stderr,("Error converting '%s' or '%s' to type 'int'" %(targetaddress,lineaddress))
					direction = 0
					dirstr = "?"
				
				self.value[EACH] += " %s %s \t(%s) \n" % (cmd,dirstr,name)
				#print >>sys.stderr, ("tagPLT::jmp::END")
				#print >>sys.stderr,("'%s','%s'"%(lineaddress,targetaddress))
				
			elif (memloc):
				#sys.stderr.write("tagPLT::memloc\n")
				memspot = thisline.find(memloc[0])
				mem = thisline[memspot:]
				for m in mem.split(','):
					if (m.find('(') < 0 and m[0] != "%"):
						_repr = self.getRepr(m, '()')
						_repr = sre.sub("\n", '\\\\n', _repr)
						#print >>sys.stderr,("Memloc found: %s:%s"%(m,_repr))
						self.value[EACH] += "\t\t %s" % _repr
				#print >>sys.stderr, ("tagPLT::memloc::END")

	def printTXT(self, outfile=sys.stdout):
		print >>outfile, (self)
		
