pure virtual support addedHEADmaster

1 files changed, 45 insertions, 18 deletions

diff --git a/classy.py b/classy.py
index 6d23b7b..1af146a 100644
--- a/classy.py
+++ b/classy.py
@@ -194,7 +194,9 @@ class Class(object):
 	def find_vtable_end(self):
 		check = self.vtable + 4
 
-		while Dword(check + 4) != 0:
+		while True:
+			value = Dword(check + 4)
+			if value < 0x80000000 or value > 0x817FFFFF: break
 			check += 4
 
 		print('Inferred vtable as %08x .. %08x' % (self.vtable, check))
@@ -307,6 +309,7 @@ class Method(object):
 	
 	def is_null(self): return False
 	def is_virtual(self): return False
+	def has_code(self): return True
 
 	def rename(self, name):
 		self.name = name
@@ -315,22 +318,23 @@ class Method(object):
 	def refresh(self):
 		MakeName(self.ea, "%s__%s%s" % (self.name, self.owner.mangled_name, self.mangled_args))
 
-		# figure out if any virtual xrefs include us
-		# a bit slow, but who cares
-		call_dict = db.virtual_calls
-		func = idaapi.get_func(self.ea)
-		check = func.startEA
-		end = func.endEA
+		if self.has_code():
+			# figure out if any virtual xrefs include us
+			# a bit slow, but who cares
+			call_dict = db.virtual_calls
+			func = idaapi.get_func(self.ea)
+			check = func.startEA
+			end = func.endEA
 
-		done_already = {}
+			done_already = {}
 
-		while check <= end:
-			if check in call_dict:
-				method = call_dict[check]
-				if method not in done_already:
-					method.refresh_comment()
-					done_already[method] = True
-			check += 4
+			while check <= end:
+				if check in call_dict:
+					method = call_dict[check]
+					if method not in done_already:
+						method.refresh_comment()
+						done_already[method] = True
+				check += 4
 	
 	def unlink(self, orphan=False):
 		if self.name == '__ct':
@@ -437,13 +441,20 @@ class VirtualMethod(Method):
 		SetFunctionCmt(self.ea, self.get_comment(), 0)
 
 
+class PureVirtualMethod(VirtualMethod):
+	def has_code(self): return False
+
+	def __init__(self, owner, name, vt_offset):
+		VirtualMethod.__init__(self, owner, name, owner.vtable+vt_offset, vt_offset)
+	
+	def refresh_comment(self):
+		MakeComm(self.ea, self.get_comment())
+
+
 class OverrideMethod(VirtualMethod):
 	def __init__(self, owner, name, ea, vt_offset, base):
 		VirtualMethod.__init__(self, owner, name, ea, vt_offset)
 
-		if ea == BADADDR || ea == 0:
-			raise ValueError("An override method can't be pure!!")
-
 		self.base = base
 		self.original = (hasattr(base, 'original') and base.original) or base
 		self.copy_signature(self.original)
@@ -547,6 +558,10 @@ def get_current_function():
 		func = idaapi.get_func(Dword(ea))
 	
 	if func is None:
+		# otherwise, might be entered into the method DB (as a pure virtual?)
+		if ea in db.known_methods:
+			return ea
+
 		# and if it's still none ...
 		Warning('Place the cursor on top of a function.')
 		raise ValueError
@@ -728,7 +743,19 @@ def register_virtual_method():
 		Warning('This class does not have a virtual table defined.')
 		return
 
+	# check for pure virtuals and handle them separately
+	current_ea = idaapi.get_screen_ea()
+	if Dword(current_ea) == 0 and current_ea > cls.vtable and current_ea <= cls.vtable_end:
+		offset = current_ea - cls.vtable
+
+		name = idaapi.askstr(idaapi.HIST_IDENT, 'vf%02X' % offset, 'Enter a name for this pure virtual')
+		method = PureVirtualMethod(cls, name, offset)
+		method.refresh()
+		Refresh()
+		return
+
 	ea = get_current_function()
+
 	unlink_method_if_exists(ea)
 
 	# try to find it within the vtable