API tools faq
paste
Advertisement
Guest User

Untitled

a guest
Oct 15th, 2019
108
0
Never
Add comment
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
text 11.24 KB | None | 0 0
raw download clone embed print report
  1. import idautils
  2. import idc
  3. import ida_bytes
  4. from pprint import pprint, pformat
  5.  
  6.  
  7. def demangle_name(name):
  8. return idc.Demangle(name, idc.GetLongPrm(idc.INF_SHORT_DN)) or name
  9.  
  10. def get_primary_vtables():
  11. vtables = {}
  12. for ea, name in idautils.Names():
  13. dname = idc.Demangle(name, idc.GetLongPrm(idc.INF_SHORT_DN))
  14. if dname != None and 'table' in dname:
  15. if not name.endswith('_0'): # Skip secondary vtables with the _0 suffix.
  16. vtables[dname] = ea
  17. return vtables
  18.  
  19. def get_new_size_arg(ea, typename=""):
  20. cur_ea = ea
  21.  
  22. # Search backwards for a call to new
  23. found_new_call = False
  24. for i in range(50):
  25. cur_ea = idc.PrevHead(cur_ea)
  26. if idc.GetMnem(cur_ea) == "call":
  27. dname = demangle_name(idc.GetOpnd(cur_ea, 0))
  28. if dname == "operator new(unsigned __int64)":
  29. found_new_call = True
  30. break
  31.  
  32. if not found_new_call:
  33. return -1
  34.  
  35. # Found the new call, now search for the (r||e)cx register argument.
  36. for i in range(20):
  37. cur_ea = idc.PrevHead(cur_ea)
  38. if idc.GetMnem(cur_ea) == "mov" and idc.GetOpnd(cur_ea, 0).endswith('cx'):
  39. if idc.GetOpType(cur_ea, 1) == idc.o_imm:
  40. return idc.GetOperandValue(cur_ea, 1)
  41. else:
  42. print("[CWT!] ERROR: got non-immediate operand for `new` argument. Can't parse this.")
  43. print("[CWT!] \tAt:{:X}, for starting ea: {:X}, type:{}".format(cur_ea, ea, typename))
  44. return -1
  45.  
  46.  
  47. return -1
  48.  
  49.  
  50. def get_new_size_arg_from_first_xref_ea(ea, typename=""):
  51. cur_ea = -1
  52.  
  53. # Get ea of first xref.
  54. for xref in idautils.XrefsTo(ea, 0):
  55. cur_ea = xref.frm
  56. break
  57.  
  58. if cur_ea == -1:
  59. return -1
  60.  
  61. return get_new_size_arg(cur_ea, typename)
  62.  
  63. def get_free_size_arg_from_vtable_ea(vtable_ea, typename):
  64. # Read the dtor from vtable[0].
  65. dtor_ea = ida_bytes.get_qword(vtable_ea)
  66. cur_ea = dtor_ea
  67.  
  68. # Search forwards for a call to free.
  69. found_free_call = False
  70. func_end_ea = idc.GetFunctionAttr(dtor_ea, idc.FUNCATTR_END)
  71. while cur_ea < func_end_ea:
  72. cur_ea = idc.NextHead(cur_ea)
  73. if idc.GetMnem(cur_ea) == "call":
  74. dname = demangle_name(idc.GetOpnd(cur_ea, 0))
  75. if 'j_j_free' in dname:
  76. found_free_call = True
  77. break
  78.  
  79. if not found_free_call:
  80. return -1
  81.  
  82. # Now go back and find the size argument passed in (r||e)dx
  83. while cur_ea > dtor_ea:
  84. cur_ea = idc.PrevHead(cur_ea)
  85. if idc.GetMnem(cur_ea) == "mov" and idc.GetOpnd(cur_ea, 0).endswith('dx'):
  86. if idc.GetOpType(cur_ea, 1) == idc.o_imm:
  87. return idc.GetOperandValue(cur_ea, 1)
  88. else:
  89. print("[CWT_MARKERS!] ERROR: got non-immediate operand for `free` argument. Can't parse this.")
  90. print("[CWT_MARKERS!] \tAt:{:X}, for dtor: {:X}, type:{}".format(cur_ea, dtor_ea, typename))
  91. return -1
  92.  
  93. return -1
  94.  
  95. def get_exception_size(typename):
  96. # Go over each name until we get the RTTI type descriptor,
  97. # then go over it's xrefs until we find a comment in some RTTI that ida parses, but can't demangle the name of,
  98. # then grab the parsed size.
  99. for ea, name in idautils.Names():
  100. dname = idc.Demangle(name, idc.GetLongPrm(idc.INF_SHORT_DN))
  101. if dname != None and typename in dname and 'RTTI Type Descriptor' in dname:
  102. for xref in idautils.XrefsTo(ea, 0):
  103. size_ea = xref.frm+16
  104. if(Comment(size_ea) == "size of thrown object"):
  105. return ida_bytes.get_dword(size_ea)
  106.  
  107. return -1
  108.  
  109. # Goes over a function and returns a list of called function ea values.
  110. def get_called_funcs(ea):
  111. cur_ea = idc.GetFunctionAttr(ea, idc.FUNCATTR_START)
  112. end_ea = idc.GetFunctionAttr(ea, idc.FUNCATTR_END)
  113.  
  114. called_funcs = []
  115. while cur_ea < end_ea:
  116. cur_ea = idc.NextHead(cur_ea)
  117. if idc.GetMnem(cur_ea) == "call":
  118. called_funcs.append(idc.GetOperandValue(cur_ea, 0))
  119.  
  120. return called_funcs
  121.  
  122. # Tracks value moved into (r||e)cx from an lea displacement offset.
  123. def get_tracked_ctor_cx_lea_offset(ea, ctor_typename=""):
  124. tracked_reg = "rcx"
  125. func_start_ea = idc.GetFunctionAttr(ea, idc.FUNCATTR_START)
  126. cur_ea = ea
  127. while cur_ea > func_start_ea:
  128. cur_ea = idc.PrevHead(cur_ea)
  129.  
  130. mnem = idc.GetMnem(cur_ea)
  131. if mnem == "mov" and idc.GetOpnd(cur_ea, 0) == tracked_reg:
  132. if idc.GetOpType(cur_ea, 1) == idc.o_reg:
  133. tracked_reg = idc.GetOpnd(cur_ea, 1)
  134. elif idc.GetOpType(cur_ea, 1) == idc.o_displ:
  135. return idc.GetOperandValue(cur_ea, 1)
  136. else:
  137. print("[CWT_MARKERS!] Got non reg -> reg mov while tracking ctor rcx lea at 0x{:X} for ctor: 0x{:X}({})! Can't parse!".format(cur_ea, ea, ctor_typename))
  138. return -1
  139.  
  140. elif mnem == "lea":
  141. if idc.GetOpnd(cur_ea, 0) == tracked_reg:
  142. if idc.GetOpType(cur_ea, 1) == idc.o_displ:
  143. return idc.GetOperandValue(cur_ea, 1)
  144.  
  145. elif mnem == "add":
  146. if idc.GetOpnd(cur_ea, 0) == tracked_reg:
  147. if idc.GetOpType(cur_ea, 1) == idc.o_imm:
  148. return idc.GetOperandValue(cur_ea, 1)
  149.  
  150.  
  151.  
  152. return 0
  153.  
  154. def get_string_ea(s):
  155. return idc.FindBinary(idaapi.get_imagebase(), idc.SEARCH_NEXT|idc.SEARCH_DOWN|idc.SEARCH_CASE, '"{}"'.format(s))
  156.  
  157. def method_ea_from_xrefed_str(s):
  158. return idc.GetFunctionAttr(list(idautils.XrefsTo(get_string_ea(s)))[0].frm, idc.FUNCATTR_START)
  159.  
  160. def method_ea_from_aob(aob):
  161. return id
  162.  
  163. def main():
  164. type_to_info = {}
  165.  
  166. # Parse out the type name from the vftable symbol.
  167. for key, value in get_primary_vtables().iteritems():
  168. typename = key
  169. if typename.endswith("::`vftable'"):
  170. typename = typename[:-11]
  171. if typename.startswith("const "):
  172. typename = typename[6:]
  173.  
  174. type_to_info[typename] = {'vtable_ea' :value, 'ctor_ea': -1, 'ctor_inlined': False, 'fields': []}
  175.  
  176. # Get a list of dtor-related functions so that we don't get a dtor when trying to xref for a ctor.
  177. dtor_related_functions = []
  178. for typename, typeinfo in type_to_info.iteritems():
  179. dtor_ea = ida_bytes.get_qword(typeinfo['vtable_ea'])
  180. dtor_related_functions.append(dtor_ea)
  181.  
  182. dtor_called = get_called_funcs(dtor_ea)
  183. dtor_related_functions.extend(dtor_called)
  184.  
  185. # Find out which ctor belongs to which class
  186. # by checking which vtable is refenced first in each.
  187. ctor_functions = {}
  188. for typename, typeinfo in type_to_info.iteritems():
  189. for xref in idautils.XrefsTo(typeinfo['vtable_ea'], 0):
  190. ctor_ea = idc.GetFunctionAttr(xref.frm, idc.FUNCATTR_START)
  191. if ctor_ea not in dtor_related_functions:
  192. if ctor_ea not in ctor_functions:
  193. ctor_functions[ctor_ea] = []
  194.  
  195. ctor_functions[ctor_ea].append((xref.frm, typename))
  196. break
  197.  
  198.  
  199. for ctor_ea in ctor_functions.keys():
  200. # Find the owner
  201. ctor_owner = (0xFFFFFFFFFFFFFFFF, '')
  202. for vtable_ref in ctor_functions[ctor_ea]:
  203. if vtable_ref[0] < ctor_owner[0]:
  204. ctor_owner = vtable_ref
  205.  
  206. # Set the class's ctor_ea's
  207. for vtable_ref in ctor_functions[ctor_ea]:
  208. if vtable_ref[1] == ctor_owner[1]:
  209. type_to_info[vtable_ref[1]]['ctor_ea'] = ctor_ea
  210. else:
  211. type_to_info[vtable_ref[1]]['ctor_inlined'] = True
  212.  
  213.  
  214. # Try to get class size.
  215. for typename, typeinfo in type_to_info.iteritems():
  216. # First try to the size via `free` in vtable[0] dtor.
  217. size = get_free_size_arg_from_vtable_ea(typeinfo['vtable_ea'], typename)
  218.  
  219. # If the size wasn't gotten via `free`, try to get it via the `new` call before ctor.
  220. if size == -1 and not type_to_info[typename]['ctor_inlined']:
  221. size = get_new_size_arg_from_first_xref_ea(typeinfo['ctor_ea'], typename)
  222.  
  223. # If the size still wasn't gotten, the ctor might be inlined:
  224. if size == -1:
  225. size = get_new_size_arg_from_first_xref_ea(typeinfo['vtable_ea'])
  226. if size != -1:
  227. type_to_info[typename]['ctor_inlined'] = True
  228.  
  229. # Or maybe it's an exception after all.
  230. if size == -1:
  231. size = get_exception_size(typename)
  232.  
  233. if size == -1:
  234. # Aww.
  235. print("[CWT_MARKERS!] Failed to get size for {}".format(typename))
  236.  
  237. type_to_info[typename]['size'] = size
  238.  
  239.  
  240.  
  241. # Make an ctor ea -> typename map for quick lookup.
  242. ctor_to_typename = {}
  243. for typename, typeinfo in type_to_info.iteritems():
  244. ctor_to_typename[typeinfo['ctor_ea']] = typename
  245.  
  246. func_to_datatype = {
  247. idc.get_name_ea(0, "InitializeCriticalSectionAndSpinCount"): "CRITICAL_SECTION",
  248. method_ea_from_xrefed_str("SRAND initialize"): "cube::World*", # cube::World::SeedExtend
  249. #method_ea_from_aob("f3 0f ?? ?? ?? ?? ?? ?? 75 ?? 80 ?? ?? ?? 00 00 01 75"): "cube::Creature*"
  250. }
  251.  
  252. # Try to get some offsets from class ctor calls and other functions.
  253. for typename, typeinfo in type_to_info.iteritems():
  254. ctor_start_ea = idc.GetFunctionAttr(typeinfo['ctor_ea'], idc.FUNCATTR_START)
  255. ctor_end_ea = idc.GetFunctionAttr(typeinfo['ctor_ea'], idc.FUNCATTR_END)
  256.  
  257. cur_ea = ctor_start_ea
  258. while cur_ea < ctor_end_ea:
  259. cur_ea = idc.NextHead(cur_ea)
  260. if idc.GetMnem(cur_ea) == "call":
  261.  
  262. # Check if it's calling a class ctor.
  263. op_val = idc.GetOperandValue(cur_ea, 0)
  264. if op_val in ctor_to_typename:
  265. member_type = ctor_to_typename[op_val]
  266. dname = demangle_name(idc.GetOpnd(cur_ea, 0))
  267.  
  268. # Make sure it's not inlined or `new`ed.
  269. if not type_to_info[member_type]['ctor_inlined']:
  270. if type_to_info[member_type]['size'] != get_new_size_arg(cur_ea):
  271.  
  272. #print("{}::ctor -> {} ({})".format(typename, dname, ctor_to_typename[op_val]))
  273. type_to_info[typename]['fields'].append((get_tracked_ctor_cx_lea_offset(cur_ea, member_type), ctor_to_typename[op_val]))
  274.  
  275. # Check if it's calling one of our other func -> datatype functions.
  276. elif op_val in func_to_datatype:
  277. datatype = func_to_datatype[op_val]
  278. type_to_info[typename]['fields'].append((get_tracked_ctor_cx_lea_offset(cur_ea, datatype), datatype))
  279.  
  280.  
  281.  
  282. # Print the classes and their sizes
  283. for typename in sorted(type_to_info):
  284. typeinfo = type_to_info[typename]
  285. if typename.startswith("cube") or typename.startswith("plasma") or typename.startswith("gfx"):
  286. print("class {} // VTable:0x{:X}, Size:0x{:X}, ctor_inlined:{}".format(typename, typeinfo['vtable_ea'], typeinfo['size'], typeinfo['ctor_inlined']))
  287. print("{")
  288. for (offset, type) in typeinfo['fields']:
  289. print("\t+{:X} {}".format(offset, type))
  290. print("};")
  291. #print("Class: {}, {} Size: 0x{:X}".format(typename, name_pad, typeinfo['size']))
  292.  
  293. #pprint(type_to_info)
  294.  
  295.  
  296.  
  297.  
  298. main()
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement
Public Pastes
Advertisement
We use cookies for various purposes including analytics. By continuing to use Pastebin, you agree to our use of cookies as described in the Cookies Policy.  OK, I Understand
Not a member of Pastebin yet?
Sign Up, it unlocks many cool features!