def get_constant_id(enumid, value): """ Return id of constant for specific value in enum. """ constid = idc.GetConstEx(enumid, value, 0, -1) if constid != idaapi.BADADDR and not idc.IsBitfield(enumid): return constid for bm in get_bitmasks(enumid): constid = idc.GetConstEx(enumid, value, 0, bm) if constid != idaapi.BADADDR: return constid return idaapi.BADADDR
def enum_member_iterate_const(enum_id): const_value = idc.GetFirstConst(enum_id, -1) while const_value != idc.BADADDR: serial = 0 const_id = idc.GetConstEx(enum_id, const_value, serial, -1) while const_id != idc.BADADDR: yield (const_id, const_value, None) serial += 1 const_id = idc.GetConstEx(enum_id, const_value, serial, -1) const_value = idc.GetNextConst(enum_id, const_value, -1)
def yacheck_enums(self): values = yaunit.load('enums') for flag in flags: for prefix, enum_width, is_bitfield, num_fields in tests: name, ea, values = str(values[0][0]), values[0][1], values[1:] eid = idc.GetEnum(name) self.assertNotEqual(eid, idaapi.BADADDR) self.assertEqual(idc.GetEnumFlag(eid), flag) self.assertEqual(idc.GetEnumName(eid), name) self.assertEqual(idc.IsBitfield(eid), is_bitfield) self.assertEqual(idc.GetEnumCmt(eid, False), prefix + 'cmt') self.assertEqual(idc.GetEnumCmt(eid, True), prefix + 'rpt') if enum_width != 0: self.assertEqual(idc.GetEnumWidth(eid), enum_width) n = 0 for value, bmask in walk_enum(eid): self.assertLessEqual(n, num_fields) v = 1 << n if is_bitfield else n self.assertEqual(value, value, v) cid = idc.GetConstEx(eid, v, 0, bmask) self.assertNotEqual(cid, idaapi.BADADDR) field = '%s_%d' % (name, n) self.assertEqual(idc.GetConstName(cid), field) # FIXME comments are not working #self.assertEqual(idc.GetConstCmt(cid, False), field + 'cmt') #self.assertEqual(idc.GetConstCmt(cid, True), field + 'rpt') n += 1 self.assertEqual(n, num_fields) if ea != None: self.assertEqual(idaapi.get_enum_id(ea, 1)[0], eid)
def make_enum_member(self, object_version, value): name = object_version.get_name() object_id = object_version.get_id() enum_object_id = object_version.get_parent_object_id() enum_id = 0 try: enum_id = self.enum_ids[enum_object_id] except: return self.hash_provider.put_hash_enum_member(idc.GetEnumName(enum_id), name, value, object_id) # create member if not idc.IsBitfield(enum_id): bmask = -1 else: bmask = object_version.get_object_flags() member_id = idc.GetConstEx(enum_id, value, 0, bmask) if member_id == idc.BADADDR: idc.AddConstEx(enum_id, name, value, bmask) member_id = idc.GetConstEx(enum_id, value, 0, bmask) else: if idc.SetConstName(member_id, name) == 0: logger.error("Failed to set const name for enum_id") # apply comments try: repeatable_headercomment = self.sanitize_comment_to_ascii( object_version.get_header_comment(True)) idc.SetConstCmt(member_id, repeatable_headercomment, 1) except KeyError: pass try: nonrepeatable_headercomment = self.sanitize_comment_to_ascii( object_version.get_header_comment(True)) idc.SetConstCmt(member_id, nonrepeatable_headercomment, 0) except KeyError: pass self.enum_member_ids[(idc.GetEnumName(enum_id), name, value)] = (member_id, object_id)
def get_const_id(eid, val, bmask): cid = idaapi.BADADDR serial = 0 while cid == idaapi.BADADDR: cid = idc.GetConstEx(eid, val, serial, bmask) serial += 1 return cid
def rename_constant(arg_ea, fct_name, arg_name, arg_enums): """ Rename constants to values from standard enumerations. """ instruction = idc.GetMnem(arg_ea) if instruction == 'push': op_num = 0 elif instruction == 'mov': op_num = 1 else: raise RenamingException('Constant: unhandled instruction ' + instruction) op_val = idc.GetOperandValue(arg_ea, op_num) # NULL if op_val == 0: targetid = idc.GetConstByName('NULL_{}_{}'.format(arg_name, fct_name)) serial = 0 enumid = idc.GetEnum(NULL_ENUM_NAME) constid = idc.GetConstEx(enumid, 0, serial, -1) while constid != idaapi.BADADDR: if constid == targetid: idc.OpEnumEx(arg_ea, op_num, enumid, serial) return serial = serial + 1 constid = idc.GetConstEx(enumid, 0, serial, -1) # All other constants op_type = idc.GetOpType(arg_ea, op_num) if op_type == idaapi.o_imm: # only one choice if len(arg_enums) == 1: enumid = idc.GetEnum(arg_enums[0]) idc.OpEnumEx(arg_ea, op_num, enumid, 0) return for enum in arg_enums: enumid = idc.GetEnum(enum) constid = get_constant_id(enumid, op_val) if constid == idaapi.BADADDR: # Not in this enum continue else: # Found the right enum idc.OpEnumEx(arg_ea, op_num, enumid, 0) return
def enum_member_iterate_bitfield(enum_id): # bitfield bmask = idc.GetFirstBmask(enum_id) while bmask != idc.BADADDR: const_value = idc.GetFirstConst(enum_id, bmask) while const_value != idc.BADADDR: # TODO must implement serial for bitfield const_id = idc.GetConstEx(enum_id, const_value, 0, bmask) yield (const_id, const_value, bmask) const_value = idc.GetNextConst(enum_id, const_value, bmask) bmask = idc.GetNextBmask(enum_id, bmask)
def enum_member_iterate_all(enum_id): const_value = idc.GetFirstConst(enum_id, -1) while const_value != idc.BADADDR: serial = 0 const_id = idc.GetConstEx(enum_id, const_value, serial, -1) while const_id != idc.BADADDR: yield (const_id, const_value, idc.BADADDR) serial += 1 const_id = idc.GetConstEx(enum_id, const_value, serial, -1) const_value = idc.GetNextConst(enum_id, const_value, -1) enum_member_iterate_bitfield(enum_id) bmask = idc.GetFirstBmask(enum_id) while bmask != idc.BADADDR: const_value = idc.GetFirstConst(enum_id, bmask) while const_value != idc.BADADDR: # TODO must implement serial for bitfield const_id = idc.GetConstEx(enum_id, const_value, 0, bmask) yield (const_id, const_value, bmask) const_value = idc.GetNextConst(enum_id, const_value, bmask) bmask = idc.GetNextBmask(enum_id, bmask)
def getConsts(self, reload=False): if self._vals and not reload: return self._vals res = {} bm = idc.GetFirstBmask(self.id) while True: if bm not in res: res[bm] = {} c = idc.GetFirstConst(self.id, bm) while c != idc.BADADDR: cid = idc.GetConstEx(self.id, c, 0, bm) res[bm][c] = cid c = idc.GetNextConst(self.id, c, bm) if bm == idc.BADADDR: self._vals = res return res bm = idc.GetNextBmask(self.id, bm)
def fix_names(self): """ Fix the table of imports and map enums to apis. """ start_addr = idc.AskAddr(idc.here(), "Enter table start address") ## check if address is within the base address and maximum address if (start_addr < idc.MinEA()) or (start_addr > idc.MaxEA()): idc.Warning("You have entered an invalid start address") idc.Exit self.start_addr = start_addr current_addr = self.start_addr #Current size of PoisonIvy IAT end_addr = current_addr + 568 # Walk the table 8 bytes at a time while current_addr <= end_addr: idc.MakeQword(current_addr) print "DEBUG: Current address - 0x%08x" % current_addr addr = idc.Qword(current_addr) print "DEBUG: address - 0x%08x" % addr if addr == -1: print "[!] Skipping address 0x%08x - 0x%08x" % (current_addr, addr) current_addr += 8 continue # Make the current address an offset idc.OpOff(current_addr,0,0) # We need to undefine the bytes incase IDA autoanalysis had converted an incorrect byte idc.MakeUnkn(addr, 1) # Create code at this address idc.MakeCode(addr) # Create function at the same address idc.MakeFunction(addr, addr+16) # Read the second operand at the address which should be the negative API address value imp_addr = idc.GetOperandValue(addr, 1) if imp_addr == -1: print "[!] Couldn't get operand at address - 0x%08x" % addr current_addr +=8 continue # try: # int_addr = int(imp_addr,16) # except ValueError as e: # print "[!] Failed on: %s - %s\n" % (imp_addr, e) # current_addr +=8 # continue # if we know about this value then let's do the work if imp_addr in self.enums.keys(): enum_id = self.enums[imp_addr] # Convert operand to enum idc.OpEnumEx(addr, 1, enum_id,0) const_id = idc.GetConstEx(enum_id, imp_addr, 0, -1) fn_name = "fn_"+idc.GetConstName(const_id) off_name = "ptr_"+idc.GetConstName(const_id) # Rename the function to the symbol name. # We append fn_ to the symbol for the function name # and ptr_ to the offset in the table. if not idc.MakeNameEx(addr, fn_name, idaapi.SN_NOWARN): print "[!] Failed to rename function %s at 0x%08x\n" % (fn_name, addr) if not idc.MakeNameEx(current_addr, off_name, idaapi.SN_NOWARN): print "[!] Failed to rename offset %s at 0x%08x\n" % (off_name,current_addr) current_addr += 8 return