diff --git a/idalib-sys/src/func_extras.h b/idalib-sys/src/func_extras.h index f7cbd646ac..6461f3addf 100644 --- a/idalib-sys/src/func_extras.h +++ b/idalib-sys/src/func_extras.h @@ -3,6 +3,7 @@ #include "pro.h" #include "funcs.hpp" #include "gdl.hpp" +#include "name.hpp" #include #include @@ -42,3 +43,20 @@ rust::Slice idalib_qbasic_block_succs(qbasic_block_t const *blk) { rust::Slice idalib_qbasic_block_preds(qbasic_block_t const *blk) { return rust::Slice(std::begin(blk->pred), std::size(blk->pred)); } + +bool idalib_func_set_name(const func_t *f, const char *name, int flags) { + if (f == nullptr || name == nullptr) { + return false; + } + return set_name(f->start_ea, name, flags); +} + +void idalib_func_set_noret(func_t *f, bool noret) { + if (f != nullptr) { + if (noret) { + f->flags |= FUNC_NORET; + } else { + f->flags &= ~FUNC_NORET; + } + } +} diff --git a/idalib-sys/src/idalib_extras.h b/idalib-sys/src/idalib_extras.h index 4e68354488..04a24c0586 100644 --- a/idalib-sys/src/idalib_extras.h +++ b/idalib-sys/src/idalib_extras.h @@ -1,7 +1,15 @@ #pragma once #include "idalib.hpp" +#include "name.hpp" bool idalib_get_library_version(int *major, int *minor, int *build) { return get_library_version(*major, *minor, *build); } + +bool idalib_set_name(uval_t ea, const char *name, int flags) { + if (name == nullptr) { + return false; + } + return set_name(ea, name, flags); +} diff --git a/idalib-sys/src/lib.rs b/idalib-sys/src/lib.rs index 2be40dbf75..a5ce4934bc 100644 --- a/idalib-sys/src/lib.rs +++ b/idalib-sys/src/lib.rs @@ -73,6 +73,7 @@ include_cpp! { #include "pro.h" #include "segment.hpp" #include "strlist.hpp" + #include "typeinf.hpp" #include "ua.hpp" #include "xref.hpp" @@ -123,6 +124,7 @@ include_cpp! { generate!("func_t") generate!("lock_func") generate!("get_func") + generate!("get_fchunk") generate!("get_func_num") generate!("get_func_qty") generate!("getn_func") @@ -384,6 +386,12 @@ include_cpp! { generate!("is_in_nlist") generate!("is_public_name") generate!("is_weak_name") + + // types + generate!("parse_decls") + generate!("get_idati") + generate!("get_ordinal_limit") + generate!("get_numbered_type_name") } pub mod hexrays { @@ -737,6 +745,7 @@ mod ffix { include!("segm_extras.h"); include!("search_extras.h"); include!("strings_extras.h"); + include!("types_extras.h"); type c_short = autocxx::c_short; type c_int = autocxx::c_int; @@ -783,6 +792,8 @@ mod ffix { unsafe fn idalib_func_flags(f: *const func_t) -> u64; unsafe fn idalib_func_name(f: *const func_t) -> Result; + unsafe fn idalib_func_set_name(f: *const func_t, name: *const c_char, flags: c_int) -> bool; + unsafe fn idalib_func_set_noret(f: *mut func_t, noret: bool); unsafe fn idalib_func_flow_chart( f: *mut func_t, @@ -999,6 +1010,7 @@ mod ffix { unsafe fn idalib_segm_perm(s: *const segment_t) -> u8; unsafe fn idalib_segm_bitness(s: *const segment_t) -> u8; unsafe fn idalib_segm_type(s: *const segment_t) -> u8; + unsafe fn idalib_segm_set_perm(s: *mut segment_t, perm: u8); unsafe fn idalib_get_cmt(ea: c_ulonglong, rptble: bool) -> String; @@ -1038,6 +1050,18 @@ mod ffix { minor: *mut c_int, build: *mut c_int, ) -> bool; + unsafe fn idalib_set_name(ea: c_ulonglong, name: *const c_char, flags: c_int) -> bool; + + unsafe fn idalib_parse_header_file(filename: *const c_char) -> c_int; + unsafe fn idalib_tinfo_get_name_by_ordinal(ordinal: u32) -> Result; + unsafe fn idalib_is_valid_type_ordinal(ordinal: u32) -> bool; + unsafe fn idalib_get_type_ordinal_limit() -> u32; + + // Type assignment functions + unsafe fn idalib_apply_type_by_ordinal(ea: c_ulonglong, ordinal: u32, flags: u32) -> bool; + unsafe fn idalib_apply_type_by_decl(ea: c_ulonglong, decl: *const c_char) -> bool; + unsafe fn idalib_get_type_ordinal_at_address(ea: c_ulonglong) -> u32; + unsafe fn idalib_get_type_string_at_address(ea: c_ulonglong) -> Result; } } @@ -1112,11 +1136,11 @@ pub mod insn { pub mod func { pub use super::ffi::{ - calc_thunk_func_target, fc_block_type_t, func_t, gdl_graph_t, get_func, get_func_num, + calc_thunk_func_target, fc_block_type_t, func_t, gdl_graph_t, get_func, get_fchunk, get_func_num, get_func_qty, getn_func, lock_func, qbasic_block_t, qflow_chart_t, }; pub use super::ffix::{ - idalib_func_flags, idalib_func_flow_chart, idalib_func_name, idalib_qbasic_block_preds, + idalib_func_flags, idalib_func_flow_chart, idalib_func_name, idalib_func_set_name, idalib_func_set_noret, idalib_qbasic_block_preds, idalib_qbasic_block_succs, idalib_qflow_graph_getn_block, }; @@ -1158,7 +1182,7 @@ pub mod segment { pub use super::ffix::{ idalib_segm_align, idalib_segm_bitness, idalib_segm_bytes, idalib_segm_name, - idalib_segm_perm, idalib_segm_type, + idalib_segm_perm, idalib_segm_set_perm, idalib_segm_type, }; } @@ -1233,6 +1257,7 @@ pub mod name { get_nlist_ea, get_nlist_idx, get_nlist_name, get_nlist_size, is_in_nlist, is_public_name, is_weak_name, }; + pub use super::ffix::idalib_set_name; } pub mod ida { @@ -1416,3 +1441,15 @@ pub mod ida { } } } + +pub mod types { + pub use super::ffi::{ + get_idati, get_ordinal_limit, get_numbered_type_name, + }; + pub use super::ffix::{ + idalib_get_type_ordinal_limit, idalib_parse_header_file, + idalib_tinfo_get_name_by_ordinal, idalib_is_valid_type_ordinal, + idalib_apply_type_by_ordinal, idalib_apply_type_by_decl, + idalib_get_type_ordinal_at_address, idalib_get_type_string_at_address, + }; +} diff --git a/idalib-sys/src/segm_extras.h b/idalib-sys/src/segm_extras.h index 4649bb656d..252879238f 100644 --- a/idalib-sys/src/segm_extras.h +++ b/idalib-sys/src/segm_extras.h @@ -43,3 +43,9 @@ std::uint8_t idalib_segm_perm(const segment_t *s) { std::uint8_t idalib_segm_type(const segment_t *s) { return s->type; } + +void idalib_segm_set_perm(segment_t *s, std::uint8_t perm) { + if (s != nullptr) { + s->perm = perm; + } +} diff --git a/idalib-sys/src/types_extras.h b/idalib-sys/src/types_extras.h new file mode 100644 index 0000000000..19e1cbc121 --- /dev/null +++ b/idalib-sys/src/types_extras.h @@ -0,0 +1,117 @@ +#pragma once + +#include "cxx.h" +#include "pro.h" +#include "typeinf.hpp" + +#include +#include + +// Parse types from a header file +inline int idalib_parse_header_file(const char *filename) { + if (filename == nullptr) { + return -1; + } + + til_t *til = get_idati(); + if (til == nullptr) { + return -1; + } + + // HTI_FIL = input is filename, HTI_MAC = define macros from base tils, + // HTI_NWR = no warnings + int flags = HTI_FIL | HTI_MAC | HTI_NWR; + return parse_decls(til, filename, nullptr, flags); +} + +// Get type name from tinfo_t (using void* to avoid direct tinfo_t exposure) +inline rust::String idalib_tinfo_get_name_by_ordinal(std::uint32_t ordinal) { + tinfo_t tif; + + if (!tif.get_numbered_type(get_idati(), ordinal)) { + return rust::String(); + } + + const char *name = tif.dstr(); + if (name == nullptr) { + return rust::String(); + } + + return rust::String(name); +} + +// Check if a type ordinal is valid +inline bool idalib_is_valid_type_ordinal(std::uint32_t ordinal) { + tinfo_t tif; + return tif.get_numbered_type(get_idati(), ordinal); +} + +// Get the maximum ordinal for type iteration +inline std::uint32_t idalib_get_type_ordinal_limit() { + return get_ordinal_limit(get_idati()); +} + +// Apply type to an address using ordinal +inline bool idalib_apply_type_by_ordinal(std::uint64_t ea, + std::uint32_t ordinal, + std::uint32_t flags) { + tinfo_t tif; + + if (!tif.get_numbered_type(get_idati(), ordinal)) { + return false; + } + + return apply_tinfo(ea, tif, flags); +} + +// Apply type to an address using C declaration string +inline bool idalib_apply_type_by_decl(std::uint64_t ea, const char *decl) { + if (decl == nullptr) { + return false; + } + + til_t *til = get_idati(); + if (til == nullptr) { + return false; + } + + return apply_cdecl(til, ea, decl); +} + +// Get type information at an address (returns ordinal, 0 if no type) +inline std::uint32_t idalib_get_type_ordinal_at_address(std::uint64_t ea) { + tinfo_t tif; + + if (!guess_tinfo(&tif, ea)) { + return 0; + } + + // Try to find the ordinal for this type + std::uint32_t limit = get_ordinal_limit(get_idati()); + for (std::uint32_t i = 1; i < limit; i++) { + tinfo_t check_tif; + if (check_tif.get_numbered_type(get_idati(), i)) { + if (tif.equals_to(check_tif)) { + return i; + } + } + } + + return 0; // Type not found in numbered types +} + +// Get type declaration string at an address +inline rust::String idalib_get_type_string_at_address(std::uint64_t ea) { + tinfo_t tif; + + if (!guess_tinfo(&tif, ea)) { + return rust::String(); + } + + const char *type_str = tif.dstr(); + if (type_str == nullptr) { + return rust::String(); + } + + return rust::String(type_str); +} \ No newline at end of file diff --git a/idalib/examples/types_ls.rs b/idalib/examples/types_ls.rs new file mode 100644 index 0000000000..3c69e5991c --- /dev/null +++ b/idalib/examples/types_ls.rs @@ -0,0 +1,109 @@ +use idalib::idb::IDB; + +fn main() -> anyhow::Result<()> { + println!("Trying to open IDA database..."); + + // Open IDA database + let idb = IDB::open("./tests/ls")?; + + println!("Testing types API..."); + + // Test getting type list length + let types = idb.types(); + let len = types.len(); + println!("Found {} types in database", len); + + // Test iterating through first few types + println!("\nFirst 10 types:"); + for i in 1..std::cmp::min(len, 10) { + if let Some(typ) = types.get_by_index(i as u32) { + if let Some(name) = typ.name() { + println!(" Type {} (ordinal {}): {}", i, typ.ordinal(), name); + } else { + println!(" Type {} (ordinal {}): ", i, typ.ordinal()); + } + } + } + + // Test type iterator + println!("\nTesting type iterator (first 5):"); + for (ordinal, typ) in types.iter().take(5) { + if let Some(name) = typ.name() { + println!(" Iterator type ordinal {}: {}", ordinal, name); + } else { + println!(" Iterator type ordinal {}: ", ordinal); + } + } + + // Test type assignment functionality + println!("\nTesting type assignment..."); + + // Get the first function in the database + if let Some((_, mut func)) = idb.functions().nth(0) { + let func_addr = func.start_address(); + println!("Testing with function at address: {:#x}", func_addr); + + // Check current type + if let Some(current_type) = func.get_type() { + println!("Current function type: {:?}", current_type.name()); + } else { + println!("Function has no assigned type"); + } + + // Try to find an existing type to apply to the function + // Let's look for a simple integer type or void type + if let Some(test_type) = types.get_by_index(1) { + if let Some(type_name) = test_type.name() { + println!("Attempting to apply type '{}' to function", type_name); + match func.set_type(&test_type) { + Ok(()) => println!("Successfully applied type to function"), + Err(e) => println!("Failed to apply type: {}", e), + } + + // Check if we can retrieve the type we just set + if let Some(new_type) = func.get_type() { + println!("New type after assignment: {:?}", new_type.name()); + } + } + } + } else { + println!("No functions found in database"); + } + + // Test address-based type assignment + println!("\nTesting address-based type assignment..."); + + // Get the first segment's start address for testing + if let Some((_, segment)) = idb.segments().nth(0) { + let test_address = segment.start_address(); + println!("Testing with address: {:#x}", test_address); + + // Check current type at address + if let Some(addr_type) = idb.get_type_at_address(test_address) { + println!("Current type at address: {:?}", addr_type.name()); + } else { + println!("No type assigned at address"); + } + + // Try to apply a type from our type list to the address + if let Some(test_type) = types.get_by_index(1) { + if let Some(type_name) = test_type.name() { + println!("Attempting to apply type '{}' to address", type_name); + match test_type.apply_to_address(test_address) { + Ok(()) => println!("Successfully applied type to address"), + Err(e) => println!("Failed to apply type to address: {}", e), + } + + // Check if we can retrieve the type we just set + if let Some(new_addr_type) = idb.get_type_at_address(test_address) { + println!("New type at address: {:?}", new_addr_type.name()); + } + } + } + } else { + println!("No segments found in database"); + } + + println!("\nType assignment API test completed!"); + Ok(()) +} diff --git a/idalib/src/func.rs b/idalib/src/func.rs index 9197cb05ce..9b0499fd28 100644 --- a/idalib/src/func.rs +++ b/idalib/src/func.rs @@ -1,16 +1,22 @@ +use std::ffi::CString; use std::marker::PhantomData; use std::mem; use std::pin::Pin; use std::ptr; use autocxx::moveit::Emplace; +use autocxx::c_int; use bitflags::bitflags; use cxx::UniquePtr; use crate::ffi::func::*; use crate::ffi::xref::has_external_refs; use crate::ffi::{range_t, IDAError, BADADDR}; +use crate::ffi::types::{ + idalib_get_type_ordinal_at_address, +}; use crate::idb::IDB; +use crate::types::{Type, TypeFlags}; use crate::Address; pub struct Function<'a> { @@ -152,6 +158,29 @@ bitflags! { } } +bitflags! { + #[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)] + pub struct NameFlags: i32 { + const CHECK = 0x00; + const NOCHECK = 0x01; + const PUBLIC = 0x02; + const NON_PUBLIC = 0x04; + const WEAK = 0x08; + const NON_WEAK = 0x10; + const AUTO = 0x20; + const NON_AUTO = 0x40; + const NOLIST = 0x80; + const NOWARN = 0x100; + const LOCAL = 0x200; + const IDBENC = 0x400; + const FORCE = 0x800; + const NODUMMY = 0x1000; + const DELTAIL = 0x2000; + const MULTI = 0x4000; + const MULTI_FORCE = 0x8000; + } +} + bitflags! { #[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)] pub struct FunctionCFGFlags: i32 { @@ -214,6 +243,32 @@ impl<'a> Function<'a> { } } + pub fn set_name(&mut self, name: impl AsRef) -> Result<(), IDAError> { + let c_name = CString::new(name.as_ref()).map_err(IDAError::ffi)?; + let success = unsafe { idalib_func_set_name(self.ptr, c_name.as_ptr(), c_int(0)) }; + if success { + Ok(()) + } else { + Err(IDAError::ffi_with(format!( + "failed to set function name to '{}'", + name.as_ref() + ))) + } + } + + pub fn set_name_with_flags(&mut self, name: impl AsRef, flags: NameFlags) -> Result<(), IDAError> { + let c_name = CString::new(name.as_ref()).map_err(IDAError::ffi)?; + let success = unsafe { idalib_func_set_name(self.ptr, c_name.as_ptr(), c_int(flags.bits())) }; + if success { + Ok(()) + } else { + Err(IDAError::ffi_with(format!( + "failed to set function name to '{}' with flags {:?}", + name.as_ref(), flags + ))) + } + } + pub fn flags(&self) -> FunctionFlags { let bits = unsafe { idalib_func_flags(self.ptr) }; FunctionFlags::from_bits_retain(bits) @@ -227,6 +282,10 @@ impl<'a> Function<'a> { unsafe { (*self.ptr).does_return() } } + pub fn set_noret(&mut self, noret: bool) { + unsafe { idalib_func_set_noret(self.ptr, noret) }; + } + pub fn analyzed_sp(&self) -> bool { unsafe { (*self.ptr).analyzed_sp() } } @@ -261,6 +320,28 @@ impl<'a> Function<'a> { _marker: PhantomData, }) } + + /// Get the type assigned to this function, if any + pub fn get_type(&self) -> Option { + let ordinal = unsafe { idalib_get_type_ordinal_at_address(self.start_address().into()) }; + if ordinal == 0 { + None + } else { + Some(Type::from_ordinal(ordinal)) + } + } + + + /// Apply a type to this function using a Type object + pub fn set_type(&mut self, typ: &Type) -> Result<(), IDAError> { + typ.apply_to_address(self.start_address()) + } + + /// Apply a type to this function using a Type object with specific flags + pub fn set_type_with_flags(&mut self, typ: &Type, flags: TypeFlags) -> Result<(), IDAError> { + typ.apply_to_address_with_flags(self.start_address(), flags) + } + } impl<'a> FunctionCFG<'a> { diff --git a/idalib/src/idb.rs b/idalib/src/idb.rs index 6468f3a5fe..8611c29908 100644 --- a/idalib/src/idb.rs +++ b/idalib/src/idb.rs @@ -3,27 +3,34 @@ use std::marker::PhantomData; use std::mem::MaybeUninit; use std::path::{Path, PathBuf}; +use autocxx::c_int; + use crate::ffi::BADADDR; use crate::ffi::bytes::*; use crate::ffi::comments::{append_cmt, idalib_get_cmt, set_cmt}; use crate::ffi::conversions::idalib_ea2str; use crate::ffi::entry::{get_entry, get_entry_ordinal, get_entry_qty}; -use crate::ffi::func::{get_func, get_func_qty, getn_func}; +use crate::ffi::func::{get_func, get_fchunk, get_func_qty, getn_func}; use crate::ffi::hexrays::{decompile_func, init_hexrays_plugin, term_hexrays_plugin}; use crate::ffi::ida::{ auto_wait, close_database_with, make_signatures, open_database_quiet, set_screen_ea, }; use crate::ffi::insn::decode; use crate::ffi::loader::find_plugin; +use crate::ffi::name::idalib_set_name; use crate::ffi::processor::get_ph; use crate::ffi::search::{idalib_find_defined, idalib_find_imm, idalib_find_text}; use crate::ffi::segment::{get_segm_by_name, get_segm_qty, getnseg, getseg}; +use crate::ffi::types::{ + idalib_parse_header_file, + idalib_get_type_ordinal_at_address, +}; use crate::ffi::util::{is_align_insn, next_head, prev_head, str2reg}; use crate::ffi::xref::{xrefblk_t, xrefblk_t_first_from, xrefblk_t_first_to}; use crate::bookmarks::Bookmarks; use crate::decompiler::CFunction; -use crate::func::{Function, FunctionId}; +use crate::func::{Function, FunctionId, NameFlags}; use crate::insn::{Insn, Register}; use crate::meta::{Metadata, MetadataMut}; use crate::name::NameList; @@ -31,6 +38,7 @@ use crate::plugin::Plugin; use crate::processor::Processor; use crate::segment::{Segment, SegmentId}; use crate::strings::StringList; +use crate::types::{Type, TypeList}; use crate::xref::{XRef, XRefQuery}; use crate::{Address, AddressFlags, IDAError, IDARuntimeHandle, prepare_library}; @@ -206,6 +214,16 @@ impl IDB { Some(Function::from_ptr(ptr)) } + pub fn function_containing_address(&self, ea: Address) -> Option { + let ptr = unsafe { get_fchunk(ea.into()) }; + + if ptr.is_null() { + return None; + } + + Some(Function::from_ptr(ptr)) + } + pub fn next_head(&self, ea: Address) -> Option
{ self.next_head_with(ea, BADADDR.into()) } @@ -414,6 +432,50 @@ impl IDB { } } + pub fn set_name(&mut self, ea: Address, name: impl AsRef) -> Result<(), IDAError> { + let c_name = CString::new(name.as_ref()).map_err(IDAError::ffi)?; + let success = unsafe { idalib_set_name(ea.into(), c_name.as_ptr(), c_int(0)) }; + if success { + Ok(()) + } else { + Err(IDAError::ffi_with(format!( + "failed to set name '{}' at address {ea:#x}", + name.as_ref() + ))) + } + } + + pub fn set_name_with_flags(&mut self, ea: Address, name: impl AsRef, flags: NameFlags) -> Result<(), IDAError> { + let c_name = CString::new(name.as_ref()).map_err(IDAError::ffi)?; + let success = unsafe { idalib_set_name(ea.into(), c_name.as_ptr(), c_int(flags.bits())) }; + if success { + Ok(()) + } else { + Err(IDAError::ffi_with(format!( + "failed to set name '{}' with flags {:?} at address {ea:#x}", + name.as_ref(), flags + ))) + } + } + + pub fn delete_name(&mut self, ea: Address) -> Result<(), IDAError> { + let success = unsafe { idalib_set_name(ea.into(), std::ptr::null(), c_int(0)) }; + if success { + Ok(()) + } else { + Err(IDAError::ffi_with(format!( + "failed to delete name at address {ea:#x}" + ))) + } + } + + pub fn set_function_name(&mut self, address: Address, name: impl AsRef) -> Result<(), IDAError> { + let mut function = self.function_at(address).ok_or_else(|| { + IDAError::ffi_with(format!("no function found at address {address:#x}")) + })?; + function.set_name(name) + } + pub fn bookmarks(&self) -> Bookmarks { Bookmarks::new(self) } @@ -474,6 +536,35 @@ impl IDB { NameList::new(self) } + pub fn types(&self) -> TypeList { + TypeList::new(self) + } + + pub fn parse_types_from_header>(&self, header_path: P) -> Result { + let path_str = header_path.as_ref().to_string_lossy(); + let c_path = CString::new(path_str.as_ref()).map_err(IDAError::ffi)?; + + let result = unsafe { idalib_parse_header_file(c_path.as_ptr()) }; + let result_val: i32 = result.into(); + if result_val < 0 { + Err(IDAError::ffi_with("Failed to parse header file")) + } else { + Ok(result_val) + } + } + + + /// Get the type at an address, if any + pub fn get_type_at_address(&self, address: Address) -> Option { + let ordinal = unsafe { idalib_get_type_ordinal_at_address(address.into()) }; + if ordinal == 0 { + None + } else { + Some(Type::from_ordinal(ordinal)) + } + } + + pub fn address_to_string(&self, ea: Address) -> Option { let s = unsafe { idalib_ea2str(ea.into()) }; diff --git a/idalib/src/lib.rs b/idalib/src/lib.rs index 053e65bafc..d940d78d74 100644 --- a/idalib/src/lib.rs +++ b/idalib/src/lib.rs @@ -92,6 +92,7 @@ pub mod plugin; pub mod processor; pub mod segment; pub mod strings; +pub mod types; pub mod xref; pub use idalib_sys as ffi; diff --git a/idalib/src/segment.rs b/idalib/src/segment.rs index 10e6a857ec..2651a90da0 100644 --- a/idalib/src/segment.rs +++ b/idalib/src/segment.rs @@ -263,6 +263,10 @@ impl<'a> Segment<'a> { SegmentPermissions::from_bits_retain(bits) } + pub fn set_permissions(&mut self, permissions: SegmentPermissions) { + unsafe { idalib_segm_set_perm(self.ptr, permissions.bits()) }; + } + pub fn bitness(&self) -> usize { (unsafe { idalib_segm_bitness(self.ptr) }) as usize } diff --git a/idalib/src/types.rs b/idalib/src/types.rs new file mode 100644 index 0000000000..1e6231ee07 --- /dev/null +++ b/idalib/src/types.rs @@ -0,0 +1,137 @@ +use std::marker::PhantomData; + +use crate::ffi::types::{ + idalib_apply_type_by_ordinal, idalib_get_type_ordinal_limit, idalib_is_valid_type_ordinal, + idalib_tinfo_get_name_by_ordinal, +}; +use crate::idb::IDB; +use crate::{Address, IDAError}; + +pub type TypeIndex = u32; + +/// Flags for type application +#[repr(u32)] +pub enum TypeFlags { + /// This is a guessed type + GUESSED = 0x0000, + /// This is a definite type + DEFINITE = 0x0001, + /// For delayed function creation + DELAYFUNC = 0x0002, + /// Strict type checking + STRICT = 0x0004, +} + +pub struct Type { + // We'll store the type ordinal instead of the tinfo_t directly + ordinal: TypeIndex, +} + +impl Type { + pub(crate) fn from_ordinal(ordinal: TypeIndex) -> Self { + Self { ordinal } + } + + pub fn name(&self) -> Option { + let name = unsafe { idalib_tinfo_get_name_by_ordinal(self.ordinal) }.ok()?; + if name.is_empty() { + None + } else { + Some(name) + } + } + + /// Apply this type to an address with default flags (TINFO_DEFINITE) + pub fn apply_to_address(&self, address: Address) -> Result<(), IDAError> { + self.apply_to_address_with_flags(address, TypeFlags::DEFINITE) + } + + /// Apply this type to an address with specific flags + pub fn apply_to_address_with_flags( + &self, + address: Address, + flags: TypeFlags, + ) -> Result<(), IDAError> { + let success = + unsafe { idalib_apply_type_by_ordinal(address.into(), self.ordinal, flags as u32) }; + if success { + Ok(()) + } else { + Err(IDAError::ffi_with("Failed to apply type to address")) + } + } + + /// Get the ordinal (index) of this type + pub fn ordinal(&self) -> TypeIndex { + self.ordinal + } +} + +pub struct TypeList<'a> { + _marker: PhantomData<&'a IDB>, +} + +impl<'a> TypeList<'a> { + pub(crate) fn new(_: &'a IDB) -> Self { + Self { + _marker: PhantomData, + } + } + + pub fn get_by_index(&self, index: TypeIndex) -> Option { + if index == 0 { + return None; // Ordinals start at 1 + } + + let is_valid = unsafe { idalib_is_valid_type_ordinal(index) }; + if !is_valid { + return None; + } + + Some(Type::from_ordinal(index)) + } + + pub fn len(&self) -> usize { + let limit = unsafe { idalib_get_type_ordinal_limit() }; + if limit == 0 || limit == u32::MAX { + 0 + } else { + (limit - 1) as usize // Ordinals start at 1, so count is limit - 1 + } + } + + pub fn is_empty(&self) -> bool { + self.len() == 0 + } + + pub fn iter(&self) -> TypeListIter<'_, 'a> { + TypeListIter { + type_list: self, + current_ordinal: 1, // Start at 1, not 0 + max_ordinal: unsafe { idalib_get_type_ordinal_limit() }, + } + } +} + +pub struct TypeListIter<'s, 'a> { + type_list: &'s TypeList<'a>, + current_ordinal: u32, + max_ordinal: u32, +} + +impl<'s, 'a> Iterator for TypeListIter<'s, 'a> { + type Item = (TypeIndex, Type); + + fn next(&mut self) -> Option { + while self.current_ordinal < self.max_ordinal { + let ordinal = self.current_ordinal; + self.current_ordinal += 1; + + if let Some(typ) = self.type_list.get_by_index(ordinal) { + return Some((ordinal, typ)); + } + } + + None + } +}