from collections import namedtuple, defaultdict
import operator
import idaapi
import idautils
import idc
from . import exceptions
from .code import lines
FF_TYPES = [idc.FF_BYTE, idc.FF_WORD, idc.FF_DWORD, idc.FF_QWORD, idc.FF_OWORD, ]
FF_SIZES = [1, 2, 4, 8, 16, ]
SIZE_TO_TYPE = dict(zip(FF_SIZES, FF_TYPES))
STRUCT_ERROR_MAP = {
idc.STRUC_ERROR_MEMBER_NAME:
(exceptions.SarkErrorStructMemberName, "already has member with this name (bad name)"),
idc.STRUC_ERROR_MEMBER_OFFSET:
(exceptions.SarkErrorStructMemberOffset, "already has member at this offset"),
idc.STRUC_ERROR_MEMBER_SIZE:
(exceptions.SarkErrorStructMemberSize, "bad number of bytes or bad sizeof(type)"),
idc.STRUC_ERROR_MEMBER_TINFO:
(exceptions.SarkErrorStructMemberTinfo, "bad typeid parameter"),
idc.STRUC_ERROR_MEMBER_STRUCT:
(exceptions.SarkErrorStructMemberStruct, "bad struct id (the 1st argument)"),
idc.STRUC_ERROR_MEMBER_UNIVAR:
(exceptions.SarkErrorStructMemberUnivar, "unions can't have variable sized members"),
idc.STRUC_ERROR_MEMBER_VARLAST:
(exceptions.SarkErrorStructMemberVarlast, "variable sized member should be the last member in the structure"),
}
def struct_member_error(err, sid, name, offset, size):
"""Create and format a struct member exception.
Args:
err: The error value returned from struct member creation
sid: The struct id
name: The member name
offset: Memeber offset
size: Member size
Returns:
A ``SarkErrorAddStructMemeberFailed`` derivative exception, with an
informative message.
"""
exception, msg = STRUCT_ERROR_MAP[err]
struct_name = idaapi.get_struc_name(sid)
return exception(('AddStructMember(struct="{}", member="{}", offset={}, size={}) '
'failed: {}').format(
struct_name,
name,
offset,
size,
msg
))
def create_struct(name):
"""Create a structure.
Args:
name: The structure's name
Returns:
The sturct ID
Raises:
exceptions.SarkStructAlreadyExists: A struct with the same name already exists
exceptions.SarkCreationFailed: Struct creation failed
"""
sid = idaapi.get_struc_id(name)
if sid != idaapi.BADADDR:
# The struct already exists.
raise exceptions.SarkStructAlreadyExists("A struct names {!r} already exists.".format(name))
sid = idaapi.add_struc(idaapi.BADADDR, name, 0)
if sid == idaapi.BADADDR:
raise exceptions.SarkStructCreationFailed("Struct creation failed.")
return sid
def get_struct(name):
"""Get a struct by it's name.
Args:
name: The name of the struct
Returns:
The struct's id
Raises:
exceptions.SarkStructNotFound: is the struct does not exist.
"""
sid = idaapi.get_struc_id(name)
if sid == idaapi.BADADDR:
raise exceptions.SarkStructNotFound()
return sid
def size_to_flags(size):
return SIZE_TO_TYPE[size] | idc.FF_DATA
def add_struct_member(sid, name, offset, size):
failure = idc.add_struc_member(sid, name, offset, size_to_flags(size), -1, size)
if failure:
raise struct_member_error(failure, sid, name, offset, size)
StructOffset = namedtuple("StructOffset", "offset size")
OperandRef = namedtuple("OperandRef", "ea n")
def infer_struct_offsets(start, end, reg_name):
offsets = set()
operands = []
for line in lines(start, end):
for operand in line.insn.operands:
if not operand.has_reg(reg_name):
continue
if not operand.type.has_phrase:
continue
if not operand.base:
continue
offset = operand.offset
if offset < 0:
raise exceptions.InvalidStructOffset(
"Invalid structure offset 0x{:08X}, probably negative number.".format(offset))
size = operand.size
offsets.add(StructOffset(offset, size))
operands.append(OperandRef(line.ea, operand.n))
return offsets, operands
def get_common_register(start, end):
"""Get the register most commonly used in accessing structs.
Access to is considered for every opcode that accesses memory
in an offset from a register::
mov eax, [ebx + 5]
For every access, the struct-referencing registers, in this case
`ebx`, are counted. The most used one is returned.
Args:
start: The adderss to start at
end: The address to finish at
"""
registers = defaultdict(int)
for line in lines(start, end):
insn = line.insn
for operand in insn.operands:
if not operand.type.has_phrase:
continue
if not operand.base:
continue
register_name = operand.base
registers[register_name] += 1
return max(registers.items(), key=operator.itemgetter(1))[0]
def offset_name(offset):
"""Format an offset into a name."""
return "offset_{:X}".format(offset.offset)
def set_struct_offsets(offsets, sid):
for offset in offsets:
try:
add_struct_member(sid,
offset_name(offset),
offset.offset,
offset.size)
except exceptions.SarkErrorStructMemberName:
# Get the offset of the member with the same name
existing_offset = idc.get_member_offset(sid, offset_name(offset))
if offset.offset == existing_offset:
pass
else:
raise
except exceptions.SarkErrorStructMemberOffset:
# Get the size of the member at the same offset
if offset.size == idc.get_member_size(sid, offset.offset):
# If they are the same, all is well.
pass
def create_struct_from_offsets(name, offsets):
sid = create_struct(name)
set_struct_offsets(offsets, sid)
def apply_struct(start, end, reg_name, struct_name):
offsets, operands = infer_struct_offsets(start, end, reg_name)
sid = get_struct(struct_name)
for ea, n in operands:
insn = idautils.DecodeInstruction(ea)
idc.op_stroff(insn, n, sid, 0)
def selection_has_offsets(start, end):
for line in lines(start, end):
for operand in line.insn.operands:
if not operand.type.has_phrase:
continue
if not operand.base:
continue
return True
return False
