ida_scripts

''' To use, load the `load.py` script in IDA Pro. Supports the latest IDAPython api. '''

__docformat__ = "numpy"

__pdoc__ = {
	'utils' : False
}

from . import memory, struc, module, misc, func

__all__ = ['memory', 'struc', 'func', 'module', 'misc']

'''
Provides class(es) for interacting with the debugged process's memory.
'''

try:
	import idc
	import idaapi
	import ida_dbg
	IS_64BIT = idaapi.get_inf_structure().is_64bit()
except: pass
import struct
from typing import *

from .utils import *

FMT_TO_SIZE = {
	'b' : 1,
	'h' : 2,
	'i' : 4,
	'l' : 4,
	'q' : 8,
}

STR_MAXLEN = 10000

class Pointer:
	''' Represent an address.
	The `Pointer` class offers convenient access to read/write from a memory address.

	Init a `Pointer` object with `Pointer.__init__` or the shorthand `ptr`.

	A `Pointer` object provides convenient functions to read/write certain data types (integers, string, pointer) to its address in memory.
	These methods take an optional `value` parameter. They read and return the value if `value` is `None`, write and return success/failure as a `bool` when `value` is not None.
	
	`Pointer` can be added/subtracted from another `Pointer` or `int` to yield a new `Pointer`.
	'''

	default_string_encoding: str = 'utf-8'
	''' The default string encoding used by `Pointer.string`. Defaults to `'utf-8'`. '''

	# INIT	
	def __init__(self, addr: int):
		''' Init with an address. '''
		assert type(addr) is int, 'addr can only be int!'
		self.addr = addr

	# READ/WRITE
	def read(self, size: int, offset: int = 0) -> bytes :
		''' Read bytes from `self.addr + offset` of given size. '''
		return idc.get_bytes(self.addr + offset, size, True)

	def write(self, bytes: bytes, offset: int = 0) -> bool:
		''' Write bytes to `self.addr + offset`. '''
		return ida_dbg.write_dbg_memory(self.addr + offset, bytes) == len(bytes)

	# The following methods reads when `value` is None and writes otherwise.
	def read_write_with_struct(self, format: str, value: int = None) -> Union[int, bool]:
		''' Read/Wwrite an integer. '''
		if value is None:
			return struct.unpack(format, self.read(FMT_TO_SIZE[format.lower()]))[0]
		else:
			return self.write(struct.pack(format, value))
	
	def u8(self, value: int = None) -> Union[int, bool]:
		''' Read/Write uint8. '''
		return self.read_write_with_struct('B', value)
	
	def u16(self, value: int = None) -> Union[int, bool]:
		''' Read/Write uint16. '''
		return self.read_write_with_struct('H', value)

	def u32(self, value: int = None) -> Union[int, bool]:
		''' Read/Write uint32. '''
		return self.read_write_with_struct('I', value)

	def u64(self, value: int = None) -> Union[int, bool]:
		''' Read/Write uint64. '''
		return self.read_write_with_struct('Q', value)

	def s8(self, value: int = None) -> Union[int, bool]:
		''' Read/Write int8. '''
		return self.read_write_with_struct('b', value)

	def s16(self, value: int = None) -> Union[int, bool]:
		''' Read/Write int16. '''
		return self.read_write_with_struct('h', value)

	def s32(self, value: int = None) -> Union[int, bool]:
		''' Read/Write int32. '''
		return self.read_write_with_struct('i', value)

	def s64(self, value: int = None) -> Union[int, bool]:
		''' Read/Write int64. '''
		return self.read_write_with_struct('q', value)

	def string(self, value: str = None, encoding: str = None, read_length: int = None) -> Union[str, bool]:
		''' Read/Write string, optionally specifying a custom encoding and length to read.
		
		Parameters
		----------
		value : str, default = None
			Value to be written, the function reads if it's `None`.
		encoding : str, default = None
			Encoding to be used for reading/writing, specify `None` to use `Pointer.default_string_encoding`.
		read_length : int, defualt = None
			Length to read, leave None to read until `\\0`.

		Returns
		-------
		`str` or `bool`
			Either the read string, or a bool indicating whether write succeeded.
		'''
		if not encoding:
			encoding = Pointer.default_string_encoding
		# read
		if value is None:
			if read_length is not None:
				return self.read(read_length).decode(encoding, 'backslashreplace')
			# read null-terminated
			string = b''
			while True:
				string += self.read(1, len(string))
				if string[-1] == 0 or len(string) > STR_MAXLEN: break
			return string.decode(encoding, 'backslashreplace')
		# write
		else:
			self.write(value.encode(encoding))

	def ptr(self, value: 'Pointer' = None) -> Union['Pointer', bool]:
		''' Read/Write pointer. '''
		method = self.u64 if IS_64BIT else self.u32
		# read
		if value is None:
			return Pointer(method())
		# write
		else:
			method(value.addr)

	def hexdump_str(self, len: int = 100, offset: int = 0) -> str:
		''' Generate hexdump from `self.addr + offset` of specified length. '''
		return hexdump(self.read(len, offset), start_offset=self.addr)

	def hexdump(self, len: int = 100, offset: int = 0):
		''' Print hexdump from `self.addr + offset` of specified length. '''
		print(self.hexdump_str(len, offset))

	# OPERATORS
	def __add__(self, other: Union['Pointer', int]) -> 'Pointer':
		if type(other) is Pointer:
			other = other.addr
		return Pointer(self.addr + other)

	def __sub__(self, other: Union['Pointer', int]) -> 'Pointer':
		if type(other) is Pointer:
			other = other.addr
		return Pointer(self.addr - other)

	def __eq__(self, other: Union['Pointer', int]) -> 'Pointer':
		if type(other) is Pointer:
			other = other.addr
		return self.addr == other

	# OPERATOR ALIASES
	# for chaining
	def add(self, other: Union['Pointer', int]) -> 'Pointer':
		return self + other
	
	def sub(self, other: Union['Pointer', int]) -> 'Pointer':
		return self - other

	def __repr__(self):
		return f'*{hex(self.addr)}'

def ptr(addr: int) -> Pointer:
	''' Initialise a `Pointer`. '''
	return Pointer(addr)

''' Provides class(es) for interacting with structures (definition and instances in memory). '''

try:
	import ida_typeinf
	import ida_struct
	import ida_idaapi
	import ida_bytes
	import idc
except: pass

import re
from json import dumps
from typing import *

from .memory import Pointer
from .utils import *

@auto_repr(['id', 'name', 'dtype', 'offset', 'size'], { 'offset' : hex, 'id' : hex })
class MemberT:
	''' Represent a member of a struct type. '''

	def __init__(self, member: 'ida_struct.member_t', struct: 'StrucT'):
		''' Init with an `ida_struct.member_t` object. '''
		assert type(member) is ida_struct.member_t, 'Invalid member!'
		self.member: 'ida_struct.member_t' = member
		self.struct = struct
		''' The struct type containing this member. '''

	@property
	def id(self) -> int:
		''' Member id (mid). '''
		return self.member.id

	@property
	def name(self) -> str:
		''' Member name. '''
		return ida_struct.get_member_name(self.id)

	@property
	def dtype(self) -> 'ida_typeinf.tinfo_t':
		''' Member type, represented with `ida_typeinf.tinfo_t`. '''
		tif = ida_typeinf.tinfo_t()
		ida_struct.get_member_tinfo(tif, self.member)
		return tif

	@property
	def offset(self) -> int:
		''' Member offset in struct. '''
		return self.member.soff

	@property
	def size(self) -> int:
		''' Member size in bytes. '''
		return ida_struct.get_member_size(self.member)

	@property
	def is_gap(self) -> bool:
		return str(self.dtype)[:4] in ['char', '_BYT'] and self.name.startswith('gap')

	def instance_at(self, addr: Pointer) -> Union[Pointer, int]:
		''' Get the value from the given address.
		
		This method reads and parses the value at the given address using the appropriate type.

		Parameters
		----------
		addr : `Pointer`

		Returns
		-------
		`Pointer` or `int`
			The parsed value.
		'''
		dtype = self.dtype
		# return ptr/array as Pointer
		if dtype.is_ptr_or_array():
			return addr
		# treat everthing else as int
		else:
			assert not dtype.is_floating(), 'FP not supported yet!'
			signed = dtype.is_signed()
			length = dtype.get_size() * 8
			method = f'{"s" if signed else "u"}{length}'
			assert hasattr(addr, method), f'Cannot find method {method}, it is likely not supported.'
			return getattr(addr, method)()

	def __hash__(self) -> int:
		return self.id

	def __eq__(self, o: object) -> bool:
		return hash(self) == hash(o)

	def instance_at_struct(self, struct_addr: Pointer):
		''' Get the value from a struct at the given address.
		
		Identical to `instance_at` but applies the offset of this member.

		Parameters
		----------
		struct_addr : `Pointer`
			A pointer to an instance of the struct.

		Returns
		-------
		`Pointer` or `int`
			The parsed value.
		'''
		return self.instance_at(struct_addr + self.offset)

@auto_repr(['id', 'name', 'size'], { 'id' : hex })
class StrucT:
	''' Represent a struct type. '''

	@staticmethod
	def find(name: str) -> 'StrucT':
		''' Find struct type by name. '''
		id = ida_struct.get_struc_id(name)
		return StrucT(ida_struct.get_struc(id)) if id != ida_idaapi.BADADDR else None

	@staticmethod
	def add_struc(name: str) -> 'StrucT':
		''' Create and return an empty struct. '''
		id = ida_struct.add_struc(ida_idaapi.BADADDR, name, False)
		return StrucT(ida_struct.get_struc(id)) if id != ida_idaapi.BADADDR else None

	def __init__(self, struc: 'ida_struct.struc_t'):
		''' Init with an `ida_struct.struc_t` object. '''
		assert type(struc) is ida_struct.struc_t, 'Invalid struc!'
		self.struc: ida_struct.struc_t = struc

	@property
	def id(self) -> int:
		''' Struct id. '''
		return self.struc.id

	@property
	def name(self) -> str:
		''' Struct name. '''
		return ida_struct.get_struc_name(self.id)

	@property
	def members(self) -> List[MemberT]:
		''' All member types of struct. '''
		return [MemberT(m, self) for m in self.struc.members]

	@property
	def size(self) -> int:
		''' Struct size in bytes. '''
		return ida_struct.get_struc_size(self.struc)

	def instance_at(self, addr: Pointer) -> 'StrucI':
		''' Get an instance of this struct at the given address. '''
		return StrucI(addr, self)
	
	def __getitem__(self, name: str) -> MemberT:
		''' Get a member type by name. '''
		member = ida_struct.get_member_by_name(self.struc, name)
		return MemberT(member, self) if member else None

	def member_at_offset(self, offset: int) -> MemberT:
		''' Get a member type by offset. '''
		member = ida_struct.get_member_by_id(ida_struct.get_member_id(self.struc, offset))
		return MemberT(member[0], self) if member else None

	def member_starting_at_offset(self, offset: int) -> bool:
		'''  Get a member type by offset, member must start at offset. '''
		member = self.member_at_offset(offset)
		return member if member and member.offset == offset else None

	# MUTATING METHODS
	# These methods do not check alignments.

	@staticmethod
	def create_struc(name: str) -> 'StrucT':
		''' Create and return a new struct with a given name. '''
		tid = ida_struct.add_struc(ida_idaapi.BADADDR, name, False)
		struct = ida_struct.get_struc(tid)
		return StrucT(struct)

	def add_member(self, declaration: str, offset: int) -> MemberT:
		''' Create and return a new member in this struct type.

		`declaration` should be like `TYPE_NAME NAME[SIZE]` (Array is optional).
		'''
		# parse decl
		tinfo, name = parse_declaration(declaration)
		size = tinfo.get_size()
		# create a bytes member of equivalent size
		result = ida_struct.add_struc_member(self.struc, name, offset, idc.FF_BYTE, None, size)
		assert result == 0, f'Failed to add member: {STRUC_ERROR_MEMBER_DESCRIPTIONS[result]}'
		member = self[name]
		# apply the correct tinfo
		ida_struct.set_member_tinfo(self.struc, member.member, 0, tinfo, 2) # SET_MEMTI_COMPATIBLE 2
		return member

	def add_gap(self, offset: int, size: int) -> MemberT:
		''' Create and return a new member representing a bytes gap. '''
		name = f'gap{to_hex(offset)}'
		return self.add_member(f'_BYTE {name}[{size}]', offset)

	def delete_member(self, member: MemberT) -> bool:
		''' Delete a given member. '''
		return ida_struct.del_struc_member(self.struc, member.offset)

	def add_member_auto(self, declaration: str, offset: int) -> MemberT:
		''' Create and return a new member in this struct type, automatically reworking gaps.

		`declaration` should be like `TYPE_NAME NAME[SIZE]` (Array is optional).

		A gap is recognised by:
		1. Is a single or array of 'char' or '_BYTE'.
		2. Name starting with 'gap' (all lowercase).

		If the requested offset falls within a gap of the above definition, the gap will automatically be adjusted to accomodate the new member.

		'''
		# check if offset's in a gap

		member = self.member_at_offset(offset)
		if member:
			# check if member's a gap
			# must be byte (one or array) with name starting with 'gap'
			assert member.is_gap, 'Failed to add member: offset is occupied by a non-gap member.'
			# recreate gap(s)
			gap_start = member.offset
			gap_end = member.offset + member.size
			# 1. delete existing gap
			assert self.delete_member(member), 'Failed to delete existing gap.'
			# 2. create lower gap if necessary
			if offset > gap_start:
				self.add_gap(gap_start, offset - gap_start)
			# 3. create member
			new_member = self.add_member(declaration, offset)
			# 4. create higher gap if necessary
			higher_gap_start = new_member.offset + new_member.size
			if higher_gap_start < gap_end:
				self.add_gap(higher_gap_start, gap_end - higher_gap_start)
			return new_member
		else:
			return self.add_member(declaration, offset)

	def __hash__(self) -> int:
		return self.id

	def __eq__(self, o: object) -> bool:
		return hash(self) == hash(o)


class StrucI:
	''' Represent a struct instance. '''

	def __init__(self, addr: Pointer, struc_t: StrucT):
		''' Init with a base address and a struct type. '''
		self.addr: Pointer = addr
		''' The address of this struct. '''
		self.struc_t: StrucT = struc_t
		''' The type definition of this struct. '''

	@property
	def members(self) -> Dict[str, Any]:
		''' Collect all members of this struct as a dict. '''
		members = self.struc_t.members
		return { m.name: m.instance_at_struct(self.addr) for m in members }

	def member(self, name: str, return_pointer: bool = True):
		''' Get a member's value or `Pointer` by its name. '''
		member = self.struc_t[name]
		if not member:
			return None
		return self.addr + member.offset if return_pointer else member.instance_at_struct(self.addr)

	def __getitem__(self, name: str):
		''' Get a member's value by its name. '''
		return self.member(name)

	def __setitem__(self, name: str, value) -> bool:
		''' Set a member's value by its name. '''
		member_ptr: Pointer = self.member(name)
		assert member_ptr, f'Cannot find member {name} for struct {self.struc_t}!'
		if type(value) is str:
			member_ptr.string(value)
		# treat everthing else as int
		else:
			dtype = self.struc_t[name].dtype
			assert not dtype.is_floating(), 'FP not supported yet!'
			signed = dtype.is_signed()
			length = dtype.get_size() * 8
			method = f'{"s" if signed else "u"}{length}'
			assert hasattr(member_ptr, method), f'Cannot find accessor for {method}, it is likely not supported.'
			return getattr(member_ptr, method)(value)

	def __repr__(self):
		return f'<StrucI addr={self.addr}, struc_t={self.struc_t}, data={dumps(self.members, default=str, indent=4)}>'

	def __hash__(self) -> int:
		return hash(self.struc_t) ^ self.addr

	def __eq__(self, o: object) -> bool:
		return type(o) is StrucI and self.struc_t == o.struc_t and self.addr == o.addr

def parse_declaration(declaration):
	m = re.search(r"^(\w+[ *]+)(\w+)(\[(\d+)\])?$", declaration)
	assert m, 'Member declaration should be like `TYPE_NAME NAME[SIZE]` (Array is optional)'

	type_name, field_name, _, arr_size = m.groups()
	assert not field_name[0].isdigit(), 'Bad field name'

	result = idc.parse_decl(type_name, 0)
	assert result, 'Failed to parse member type. It should be like `TYPE_NAME NAME[SIZE]` (Array is optional)'

	_, tp, fld = result
	tinfo = ida_typeinf.tinfo_t()
	tinfo.deserialize(ida_typeinf.cvar.idati, tp, fld, None)
	if arr_size:
		assert tinfo.create_array(tinfo, int(arr_size))
	return tinfo, field_name

''' Provides class(es) for interacting with functions. '''

try:
	import idautils
	import ida_funcs
	import ida_name
	import idaapi
	import ida_hexrays
	import ida_lines
	import ida_pro
except: pass
import networkx as nx
from typing import *
from .utils import *
from .struc import *

@auto_repr(['op', 'address'], { 'id' : hex, 'address' : hex })
class CItem:
	''' Represent an item inside a graph of a decompiled function. '''

	def __init__(self, item: 'ida_hexrays.citem_t', graph: nx.DiGraph):
		''' Initialise with a `ida_hexrays.citem_t` object. '''
		self.item = item
		# python can handle retain cycles
		self.graph = graph

	@property
	def id(self) -> int:
		''' Item object id. '''
		return self.item.obj_id

	@property
	def address(self) -> int:
		''' Item address. '''
		return self.item.ea

	@property
	def op(self) -> str:
		''' Name of the operation represented by this item. '''
		return ida_hexrays.get_ctype_name(self.op_id)

	@property
	def op_id(self) -> int:
		''' The operation represented by this item. '''
		return self.item.op

	@property
	def type_name(self) -> str:
		''' The name of the type associated with this item. '''
		expr: ida_hexrays.cexpr_t = self.item.cexpr
		if self.item.is_expr() and not expr.type.empty():
			tstr = expr.type._print()
			return tstr

	@property
	def accessed_struct(self) -> StrucT:
		''' Find and return the `StrucT` accessed by this memptr/memref operation. '''
		assert self.op_id in [ida_hexrays.cot_memptr, ida_hexrays.cot_memref], 'Op is not memptr or memref.'
		if self.op_id == idaapi.cot_memptr:
			struct_tinfo = self.item.cexpr.x.type.get_pointed_object()
		elif self.op_id == idaapi.cot_memref:
			struct_tinfo = self.item.cexpr.x.type
		return StrucT.find(str(struct_tinfo))

	@property
	def accessed_struct_member(self) -> MemberT:
		''' Find and return the `MemberT` accessed by this memptr/memref operation. '''
		assert self.op_id in [ida_hexrays.cot_memptr, ida_hexrays.cot_memref], 'Op is not memptr or memref.'
		struct = self.accessed_struct
		member_offset = self.item.cexpr.m
		return struct.member_at_offset(member_offset)

	@property
	def called_function(self) -> 'Func':
		''' Find and return the `Func` called by this call operation. '''
		assert self.op_id in [ida_hexrays.cot_call], 'Op is not call!'
		return Func.func_at(self.item.x.obj_ea)

	@property
	def num_value(self) -> int:
		''' Return the value corresponding to this num operation. '''
		assert self.op_id in [ida_hexrays.cot_num], 'Op is not num!'
		return self.item.n._value

	# @property
	# def is_member_defined_in_struct(self) -> bool:
	# 	''' Checks whether there exists a member in the referenced struct type starting at the offset. '''
	# 	return self.accessed_struct.member_starting_at_offset(self.item.cexpr.m) is not None

	@property
	def label(self) -> str:
		''' Short label for item. '''
	
		def get_expr_name(expr):
			name = expr.print1(None)
			name = ida_lines.tag_remove(name)
			name = ida_pro.str2user(name)
			return name
		
		op = self.item.op
		insn = self.item.cinsn
		expr = self.item.cexpr
		parts = [ida_hexrays.get_ctype_name(op)]
		if op == ida_hexrays.cot_ptr:
			parts.append(".%d" % expr.ptrsize)
		elif op == ida_hexrays.cot_memptr:
			parts.append(".%d (m=%d)" % (expr.ptrsize, expr.m))
		elif op == ida_hexrays.cot_memref:
			parts.append(" (m=%d)" % (expr.m,))
		elif op in [
				ida_hexrays.cot_obj,
				ida_hexrays.cot_var]:
			name = get_expr_name(expr)
			parts.append(".%d %s" % (expr.refwidth, name))
		elif op in [
				ida_hexrays.cot_num,
				ida_hexrays.cot_helper,
				ida_hexrays.cot_str]:
			name = get_expr_name(expr)
			parts.append(" %s" % (name,))
		elif op == ida_hexrays.cit_goto:
			parts.append(" LABEL_%d" % insn.cgoto.label_num)
		elif op == ida_hexrays.cit_asm:
			parts.append("<asm statements; unsupported ATM>")
		if self.item.is_expr() and not expr.type.empty():
			tstr = expr.type._print()
			parts.append(tstr if tstr else "?")
		return "//".join(parts)

	@property
	def children(self) -> List['CItem']:
		''' The children of this item. '''
		return list(self.graph.successors(self))
	
	@property
	def n_children(self) -> int:
		''' The number of children of this item. '''
		return len(self.children)

	@property
	def parent(self) -> 'CItem':
		''' The parent of this item. '''
		try:
			return next(self.graph.predecessors(self))
		except:
			return None

	@property
	def parent_expression(self) -> 'CItem':
		''' Find and return the parent expression. '''
		return self.parent_where(lambda parent: parent.op == ida_hexrays.cit_expr)

	def child_where(self, condition: Callable[['CItem'], bool]) -> 'CItem':
		''' Preorder dfs search for the first child meeting the condition. '''
		dfs = nx.dfs_preorder_nodes(self.graph, self)
		node = next(dfs)
		while node:
			if condition(node):
				return node
			try:
				node = next(dfs)
			except:
				break

	def children_where(self, condition: Callable[['CItem'], bool]) -> List['CItem']:
		''' Preorder dfs search for all children meeting the condition. '''
		dfs = nx.dfs_preorder_nodes(self.graph, self)
		nodes = []
		node = next(dfs)
		while node:
			if condition(node):
				nodes.append(node)
			try:
				node = next(dfs)
			except:
				break
		return nodes

	def parent_where(self, condition: Callable[['CItem'], bool]) -> 'CItem':
		parent = self.parent
		while parent:
			if condition(parent):
				return parent
			parent = parent.parent

	def subtree_search_strict(self, query: Dict[str, Any]) -> List[List[List['CItem']]]:
		''' Perform a dfs tree search with the specified query.
		
		The query is a `dict` representing a subtree to search for, recursively defined as:
		```
		<query> = {
			"op" : <str>,
			"condition" : <lambda (CItem) -> bool>,
			"children" : [<query>, ...]
		}
		```

		`op` the name of the operation for the current item
		`condition` the matching condition for the current item
		`children` nested query(ies) specifying the children to match

		All fields are optional, but at least one of `op, condition` must be present.
		The query does not have to specify the complete subtree, but each of its specified level must be complete.
		ie. If node A has children B, C, query can either not specify A's children at all or it must specify A's children as both B and C.

		Returns results as a list of `List[List[CItem]]`, the outermost list representing the unique matches, each inner `List[List[CItem]]` representing a match.
		Example of a match:

		Query: `{ op : A, children: [{ op : B}, { op : C }]}`

		A Match: `[[A], [B, C]]`

		Returns: `[[[A], [B, C]], ...]`

		'''

		results = []

		def recursive_search(item: CItem, query: Dict[str, Any], result: List[List[CItem]] = None, index: int = 0) -> bool:
			assert 'condition' in query or 'op' in query, 'Query must have at least one of condition or op!'
			nonlocal results

			if index == 0:
				result = [[]]

			match = True
			# 1. check each of the other attributes
			attrs = { 'op' }
			for attr in attrs:
				if attr in query:
					match = match and query[attr] == getattr(item, attr)
			# 2. check condition, if exists
			if match and 'condition' in query:
				match = match and query['condition'](item)

			# always try to match deeper subtrees
			if index == 0:
				for child in item.children:
					recursive_search(child, query, None, 0)

			if match:
				# this node matches, search subtree
				subtree_match = False
				# add this item to the appropriate position
				if index == len(result):
					result.append([item])
				else:
					result[index].append(item)
				# this item matches, try to match children 'strictly'
				if item.n_children == 0 or not 'children' in query:
					# reached end of search
					# success if query has on more children
					subtree_match = not 'children' in query
				else:
					# needs to search deeper
					children = item.children
					if len(children) == len(query['children']):					
						for child_query in query['children']:
							for child_item in item.children:
								if recursive_search(child_item, child_query, result, index+1):
									children.remove(child_item)
									break

						if len(children) == 0:
							# queries 1-1 matched
							subtree_match = True
				if subtree_match and index == 0:
					results.append(result)
				return subtree_match
			
			# else this node doesn't match
			return False

		recursive_search(self, query)
		return results

	def __hash__(self):
		return self.item.obj_id

	def __eq__(self, other: 'CItem'):
		return hash(self) == hash(other)

try:
	class GraphBuilder(ida_hexrays.ctree_parentee_t):
		''' Utility class used to build the decompiled items graph. '''

		def __init__(self):
			ida_hexrays.ctree_parentee_t.__init__(self)
			self.graph = nx.DiGraph()
		
		def add_item(self, item: 'ida_hexrays.citem_t', type: str):
			parent = self.parents.back()
			item = CItem(item, self.graph)
			self.graph.add_node(item, type=type)
			if parent:
				self.graph.add_edge(CItem(parent, self.graph), item)
			return 0

		def visit_insn(self, i: 'ida_hexrays.cinsn_t'):
			return self.add_item(i, 'insn')

		def visit_expr(self, e: 'ida_hexrays.cexpr_t'):
			return self.add_item(e, 'expr')
except: pass

@auto_repr(['original'])
class CFunc:
	''' Represent a decompiled function. '''

	def __init__(self, original: 'Func', decompiled: 'ida_hexrays.cfunc_t'):
		''' Initialise with the original `Func` and the decompiled `ida_hexrays.cfunc_t` object. '''
		self.original = original
		''' The `Func` object from which this object is derived. '''
		self.decompiled = decompiled
		''' The decompiled `ida_hexrays.cfunc_t` object. '''
		# build graph
		gb = GraphBuilder()
		gb.apply_to(decompiled.body, None)
		self.body: nx.DiGraph = gb.graph
		''' The graph representing the decompiled function. '''
		self.body_root: CItem = next(nx.topological_sort(self.body))
		''' The root node of the graph for this function. '''
	
	@property
	def psuedocode(self) -> str:
		''' Psuedocode of the decompiled function. '''
		lines = self.decompiled.get_pseudocode()
		return '\n'.join([ida_lines.tag_remove(l.line) for l in lines])


@auto_repr(['name', 'start', 'end', 'size'], { 'start' : hex, 'end' : hex, 'size' : hex})
class Func:
	''' Represent a function. '''

	@staticmethod
	def find(name: str) -> 'Func':
		''' Find function by name. '''
		for addr in idautils.Functions():
			if ida_name.get_ea_name(addr) == name:
				func = idaapi.get_func(addr)
				return Func(func) if func else None

	@staticmethod
	def func_at(addr: Union[Pointer, int]) -> 'Func':
		if type(addr) is Pointer:
			addr = addr.addr
		return Func(idaapi.get_func(addr))

	def __init__(self, func: 'ida_funcs.func_t'):
		''' Initialise with a `ida_funcs.func_t` object.'''
		assert type(func) is ida_funcs.func_t, 'Invalid func!'
		self.func = func

	@property
	def name(self) -> str:
		''' Function name. '''
		return ida_name.get_ea_name(self.func.start_ea)

	@property
	def start(self) -> int:
		''' Function start address. '''
		return self.func.start_ea

	@property
	def end(self) -> int:
		''' Function end address. '''
		return self.func.end_ea

	@property
	def size(self) -> int:
		''' Function size. '''
		return self.func.size()

	@property
	def tif(self) -> 'ida_typeinf.tinfo_t':
		tif = ida_typeinf.tinfo_t()
		return tif if idaapi.get_tinfo(tif, self.func.start_ea) else None

	@property
	def arguments(self) -> Dict[str, 'ida_typeinf.tinfo_t']:
		''' Return a `dict` of argument name to type info. '''
		func_data = idaapi.func_type_data_t()
		self.tif.get_func_details(func_data)
		return { arg.name:arg.type for arg in func_data }

	def decompile(self) -> CFunc:
		''' Decompile this function. '''
		err = ida_hexrays.hexrays_failure_t()
		cfunc = ida_hexrays.decompile_func(self.func, err)
		assert cfunc, f'Decompilation failed: {err}'
		return CFunc(self, cfunc)

''' Provides class(es) for interacing with modules. '''

try:
	import idc
except: pass
from typing import *

from .utils import *

@auto_repr(['name', 'addr'], { 'addr' : hex })
class Module:
	''' Represents a single module. '''

	@staticmethod
	def all_modules() -> List['Module']:
		''' Return a list of all the modules. '''
		m = idc.get_first_module()
		modules = []
		while m:
			modules.append(m)
			m = idc.get_next_module(m)
		return modules

	@staticmethod
	def find(name: str) -> 'Module':
		modules = Module.all_modules()
		for module in modules:
			if module.name.endswith(name):
				return module

	def __init__(self, addr: int):
		self.name = idc.get_module_name(addr)
		self.addr = addr

''' Provides additional helper functions. '''

try:
	import ida_frame
	import ida_struct
	import idaapi
	import idc
except: pass

from .memory import Pointer

def find_local_var(var_name: str) -> Pointer:
	''' Find a local variable by name when paused in a function frame. '''
	frame = ida_frame.get_frame(idc.here())
	loc_var = ida_struct.get_member_by_name(frame, var_name)
	if (loc_var is None):
		return None
	stack_ptr = idc.GetRegValue('rsp' if idaapi.get_inf_structure().is_64bit() else 'esp')
	ea = loc_var.soff + stack_ptr
	return Pointer(ea)

def find_symbol(name: str) -> Pointer:
	''' Find a global symbol by name. '''
	addr = idc.get_name_ea_simple(name)
	if addr == idc.BADADDR:
		return None
	return Pointer(addr)