idalib/
idb.rs

1use std::ffi::CString;
2use std::marker::PhantomData;
3use std::mem::MaybeUninit;
4use std::path::{Path, PathBuf};
5
6use crate::ffi::bytes::*;
7use crate::ffi::comments::{append_cmt, idalib_get_cmt, set_cmt};
8use crate::ffi::conversions::idalib_ea2str;
9use crate::ffi::entry::{get_entry, get_entry_ordinal, get_entry_qty};
10use crate::ffi::func::{get_func, get_func_qty, getn_func};
11use crate::ffi::hexrays::{decompile_func, init_hexrays_plugin, term_hexrays_plugin};
12use crate::ffi::ida::{
13    auto_wait, close_database_with, make_signatures, open_database_quiet, set_screen_ea,
14};
15use crate::ffi::insn::decode;
16use crate::ffi::loader::find_plugin;
17use crate::ffi::processor::get_ph;
18use crate::ffi::search::{idalib_find_defined, idalib_find_imm, idalib_find_text};
19use crate::ffi::segment::{get_segm_qty, getnseg, getseg};
20use crate::ffi::util::{is_align_insn, next_head, prev_head, str2reg};
21use crate::ffi::xref::{xrefblk_t, xrefblk_t_first_from, xrefblk_t_first_to};
22use crate::ffi::BADADDR;
23
24use crate::bookmarks::Bookmarks;
25use crate::decompiler::CFunction;
26use crate::func::{Function, FunctionId};
27use crate::insn::{Insn, Register};
28use crate::meta::{Metadata, MetadataMut};
29use crate::plugin::Plugin;
30use crate::processor::Processor;
31use crate::segment::{Segment, SegmentId};
32use crate::strings::StringList;
33use crate::xref::{XRef, XRefQuery};
34use crate::{prepare_library, Address, IDAError, IDARuntimeHandle};
35
36pub struct IDB {
37    path: PathBuf,
38    save: bool,
39    decompiler: bool,
40    _guard: IDARuntimeHandle,
41    _marker: PhantomData<*const ()>,
42}
43
44impl IDB {
45    pub fn open(path: impl AsRef<Path>) -> Result<Self, IDAError> {
46        Self::open_with(path, true, false)
47    }
48
49    pub fn open_with(
50        path: impl AsRef<Path>,
51        auto_analyse: bool,
52        save: bool,
53    ) -> Result<Self, IDAError> {
54        let _guard = prepare_library();
55        let path = path.as_ref();
56
57        if !path.exists() || !path.is_file() {
58            return Err(IDAError::not_found(path));
59        }
60
61        open_database_quiet(path, auto_analyse)?;
62
63        let decompiler = unsafe { init_hexrays_plugin(0.into()) };
64
65        Ok(Self {
66            path: path.to_owned(),
67            save,
68            decompiler,
69            _guard,
70            _marker: PhantomData,
71        })
72    }
73
74    pub fn path(&self) -> &Path {
75        &self.path
76    }
77
78    pub fn save_on_close(&mut self, status: bool) {
79        self.save = status;
80    }
81
82    pub fn auto_wait(&mut self) -> bool {
83        unsafe { auto_wait() }
84    }
85
86    pub fn set_screen_address(&mut self, ea: Address) {
87        set_screen_ea(ea.into());
88    }
89
90    pub fn make_signatures(&mut self, only_pat: bool) -> Result<(), IDAError> {
91        make_signatures(only_pat)
92    }
93
94    pub fn decompiler_available(&self) -> bool {
95        self.decompiler
96    }
97
98    pub fn meta(&self) -> Metadata {
99        Metadata::new()
100    }
101
102    pub fn meta_mut(&mut self) -> MetadataMut {
103        MetadataMut::new()
104    }
105
106    pub fn processor(&self) -> Processor {
107        let ptr = unsafe { get_ph() };
108        Processor::from_ptr(ptr)
109    }
110
111    pub fn entries(&self) -> EntryPointIter {
112        let limit = unsafe { get_entry_qty() };
113        EntryPointIter {
114            index: 0,
115            limit,
116            _marker: PhantomData,
117        }
118    }
119
120    pub fn function_at(&self, ea: Address) -> Option<Function> {
121        let ptr = unsafe { get_func(ea.into()) };
122
123        if ptr.is_null() {
124            return None;
125        }
126
127        Some(Function::from_ptr(ptr))
128    }
129
130    pub fn next_head(&self, ea: Address) -> Option<Address> {
131        self.next_head_with(ea, BADADDR.into())
132    }
133
134    pub fn next_head_with(&self, ea: Address, max_ea: Address) -> Option<Address> {
135        let next = unsafe { next_head(ea.into(), max_ea.into()) };
136        if next == BADADDR {
137            None
138        } else {
139            Some(next.into())
140        }
141    }
142
143    pub fn prev_head(&self, ea: Address) -> Option<Address> {
144        self.prev_head_with(ea, 0)
145    }
146
147    pub fn prev_head_with(&self, ea: Address, min_ea: Address) -> Option<Address> {
148        let prev = unsafe { prev_head(ea.into(), min_ea.into()) };
149        if prev == BADADDR {
150            None
151        } else {
152            Some(prev.into())
153        }
154    }
155
156    pub fn insn_at(&self, ea: Address) -> Option<Insn> {
157        let insn = decode(ea.into())?;
158        Some(Insn::from_repr(insn))
159    }
160
161    pub fn decompile<'a>(&'a self, f: &Function<'a>) -> Result<CFunction<'a>, IDAError> {
162        self.decompile_with(f, false)
163    }
164
165    pub fn decompile_with<'a>(
166        &'a self,
167        f: &Function<'a>,
168        all_blocks: bool,
169    ) -> Result<CFunction<'a>, IDAError> {
170        if !self.decompiler {
171            return Err(IDAError::ffi_with("no decompiler available"));
172        }
173
174        Ok(unsafe {
175            decompile_func(f.as_ptr(), all_blocks)
176                .map(|f| CFunction::new(f).expect("null pointer checked"))?
177        })
178    }
179
180    pub fn function_by_id(&self, id: FunctionId) -> Option<Function> {
181        let ptr = unsafe { getn_func(id) };
182
183        if ptr.is_null() {
184            return None;
185        }
186
187        Some(Function::from_ptr(ptr))
188    }
189
190    pub fn functions<'a>(&'a self) -> impl Iterator<Item = (FunctionId, Function<'a>)> + 'a {
191        (0..self.function_count()).filter_map(|id| self.function_by_id(id).map(|f| (id, f)))
192    }
193
194    pub fn function_count(&self) -> usize {
195        unsafe { get_func_qty() }
196    }
197
198    pub fn segment_at(&self, ea: Address) -> Option<Segment> {
199        let ptr = unsafe { getseg(ea.into()) };
200
201        if ptr.is_null() {
202            return None;
203        }
204
205        Some(Segment::from_ptr(ptr))
206    }
207
208    pub fn segment_by_id(&self, id: SegmentId) -> Option<Segment> {
209        let ptr = unsafe { getnseg((id as i32).into()) };
210
211        if ptr.is_null() {
212            return None;
213        }
214
215        Some(Segment::from_ptr(ptr))
216    }
217
218    pub fn segments<'a>(&'a self) -> impl Iterator<Item = (SegmentId, Segment<'a>)> + 'a {
219        (0..self.segment_count()).filter_map(|id| self.segment_by_id(id).map(|s| (id, s)))
220    }
221
222    pub fn segment_count(&self) -> usize {
223        unsafe { get_segm_qty().0 as _ }
224    }
225
226    pub fn register_by_name(&self, name: impl AsRef<str>) -> Option<Register> {
227        let s = CString::new(name.as_ref()).ok()?;
228        let id = unsafe { str2reg(s.as_ptr()).0 };
229
230        if id == -1 {
231            None
232        } else {
233            Some(id as _)
234        }
235    }
236
237    pub fn insn_alignment_at(&self, ea: Address) -> Option<usize> {
238        let align = unsafe { is_align_insn(ea.into()).0 };
239        if align == 0 {
240            None
241        } else {
242            Some(align as _)
243        }
244    }
245
246    pub fn first_xref_from(&self, ea: Address, flags: XRefQuery) -> Option<XRef> {
247        let mut xref = MaybeUninit::<xrefblk_t>::zeroed();
248        let found =
249            unsafe { xrefblk_t_first_from(xref.as_mut_ptr(), ea.into(), flags.bits().into()) };
250
251        if found {
252            Some(XRef::from_repr(unsafe { xref.assume_init() }))
253        } else {
254            None
255        }
256    }
257
258    pub fn first_xref_to(&self, ea: Address, flags: XRefQuery) -> Option<XRef> {
259        let mut xref = MaybeUninit::<xrefblk_t>::zeroed();
260        let found =
261            unsafe { xrefblk_t_first_to(xref.as_mut_ptr(), ea.into(), flags.bits().into()) };
262
263        if found {
264            Some(XRef::from_repr(unsafe { xref.assume_init() }))
265        } else {
266            None
267        }
268    }
269
270    pub fn get_cmt(&self, ea: Address) -> Option<String> {
271        self.get_cmt_with(ea, false)
272    }
273
274    pub fn get_cmt_with(&self, ea: Address, rptble: bool) -> Option<String> {
275        let s = unsafe { idalib_get_cmt(ea.into(), rptble) };
276
277        if s.is_empty() {
278            None
279        } else {
280            Some(s)
281        }
282    }
283
284    pub fn set_cmt(&self, ea: Address, comm: impl AsRef<str>) -> Result<(), IDAError> {
285        self.set_cmt_with(ea, comm, false)
286    }
287
288    pub fn set_cmt_with(
289        &self,
290        ea: Address,
291        comm: impl AsRef<str>,
292        rptble: bool,
293    ) -> Result<(), IDAError> {
294        let s = CString::new(comm.as_ref()).map_err(IDAError::ffi)?;
295        if unsafe { set_cmt(ea.into(), s.as_ptr(), rptble) } {
296            Ok(())
297        } else {
298            Err(IDAError::ffi_with(format!(
299                "failed to set comment at {ea:#x}"
300            )))
301        }
302    }
303
304    pub fn append_cmt(&self, ea: Address, comm: impl AsRef<str>) -> Result<(), IDAError> {
305        self.append_cmt_with(ea, comm, false)
306    }
307
308    pub fn append_cmt_with(
309        &self,
310        ea: Address,
311        comm: impl AsRef<str>,
312        rptble: bool,
313    ) -> Result<(), IDAError> {
314        let s = CString::new(comm.as_ref()).map_err(IDAError::ffi)?;
315        if unsafe { append_cmt(ea.into(), s.as_ptr(), rptble) } {
316            Ok(())
317        } else {
318            Err(IDAError::ffi_with(format!(
319                "failed to append comment at {ea:#x}"
320            )))
321        }
322    }
323
324    pub fn remove_cmt(&self, ea: Address) -> Result<(), IDAError> {
325        self.remove_cmt_with(ea, false)
326    }
327
328    pub fn remove_cmt_with(&self, ea: Address, rptble: bool) -> Result<(), IDAError> {
329        let s = CString::new("").map_err(IDAError::ffi)?;
330        if unsafe { set_cmt(ea.into(), s.as_ptr(), rptble) } {
331            Ok(())
332        } else {
333            Err(IDAError::ffi_with(format!(
334                "failed to remove comment at {ea:#x}"
335            )))
336        }
337    }
338
339    pub fn bookmarks(&self) -> Bookmarks {
340        Bookmarks::new(self)
341    }
342
343    pub fn find_text(&self, start_ea: Address, text: impl AsRef<str>) -> Option<Address> {
344        let s = CString::new(text.as_ref()).ok()?;
345        let addr = unsafe { idalib_find_text(start_ea.into(), s.as_ptr()) };
346        if addr == BADADDR {
347            None
348        } else {
349            Some(addr.into())
350        }
351    }
352
353    pub fn find_text_iter<'a, T>(&'a self, text: T) -> impl Iterator<Item = Address> + 'a
354    where
355        T: AsRef<str> + 'a,
356    {
357        let mut cur = 0u64;
358        std::iter::from_fn(move || {
359            let found = self.find_text(cur, text.as_ref())?;
360            cur = self.find_defined(found).unwrap_or(BADADDR.into());
361            Some(found)
362        })
363    }
364
365    pub fn find_imm(&self, start_ea: Address, imm: u32) -> Option<Address> {
366        let addr = unsafe { idalib_find_imm(start_ea.into(), imm.into()) };
367        if addr == BADADDR {
368            None
369        } else {
370            Some(addr.into())
371        }
372    }
373
374    pub fn find_imm_iter<'a>(&'a self, imm: u32) -> impl Iterator<Item = Address> + 'a {
375        let mut cur = 0u64;
376        std::iter::from_fn(move || {
377            cur = self.find_imm(cur, imm)?;
378            Some(cur)
379        })
380    }
381
382    pub fn find_defined(&self, start_ea: Address) -> Option<Address> {
383        let addr = unsafe { idalib_find_defined(start_ea.into()) };
384        if addr == BADADDR {
385            None
386        } else {
387            Some(addr.into())
388        }
389    }
390
391    pub fn strings(&self) -> StringList {
392        StringList::new(self)
393    }
394
395    pub fn address_to_string(&self, ea: Address) -> Option<String> {
396        let s = unsafe { idalib_ea2str(ea.into()) };
397
398        if s.is_empty() {
399            None
400        } else {
401            Some(s)
402        }
403    }
404
405    pub fn get_byte(&self, ea: Address) -> u8 {
406        unsafe { idalib_get_byte(ea.into()) }
407    }
408
409    pub fn get_word(&self, ea: Address) -> u16 {
410        unsafe { idalib_get_word(ea.into()) }
411    }
412
413    pub fn get_dword(&self, ea: Address) -> u32 {
414        unsafe { idalib_get_dword(ea.into()) }
415    }
416
417    pub fn get_qword(&self, ea: Address) -> u64 {
418        unsafe { idalib_get_qword(ea.into()) }
419    }
420
421    pub fn get_bytes(&self, ea: Address, size: usize) -> Vec<u8> {
422        let mut buf = Vec::with_capacity(size);
423
424        let Ok(new_len) = (unsafe { idalib_get_bytes(ea.into(), &mut buf) }) else {
425            return Vec::with_capacity(0);
426        };
427
428        unsafe {
429            buf.set_len(new_len);
430        }
431
432        buf
433    }
434
435    pub fn find_plugin(
436        &self,
437        name: impl AsRef<str>,
438        load_if_needed: bool,
439    ) -> Result<Plugin, IDAError> {
440        let plugin = CString::new(name.as_ref()).map_err(IDAError::ffi)?;
441        let ptr = unsafe { find_plugin(plugin.as_ptr(), load_if_needed) };
442
443        if ptr.is_null() {
444            Err(IDAError::ffi_with(format!(
445                "failed to load {} plugin",
446                name.as_ref()
447            )))
448        } else {
449            Ok(Plugin::from_ptr(ptr))
450        }
451    }
452
453    pub fn load_plugin(&self, name: impl AsRef<str>) -> Result<Plugin, IDAError> {
454        self.find_plugin(name, true)
455    }
456}
457
458impl Drop for IDB {
459    fn drop(&mut self) {
460        if self.decompiler {
461            unsafe {
462                term_hexrays_plugin();
463            }
464        }
465        close_database_with(self.save);
466    }
467}
468
469pub struct EntryPointIter<'a> {
470    index: usize,
471    limit: usize,
472    _marker: PhantomData<&'a IDB>,
473}
474
475impl<'a> Iterator for EntryPointIter<'a> {
476    type Item = Address;
477
478    fn next(&mut self) -> Option<Self::Item> {
479        if self.index >= self.limit {
480            return None;
481        }
482
483        let ordinal = unsafe { get_entry_ordinal(self.index) };
484        let addr = unsafe { get_entry(ordinal) };
485
486        // skip?
487        if addr == BADADDR {
488            self.index += 1;
489            return self.next();
490        }
491
492        Some(addr.into())
493    }
494
495    fn size_hint(&self) -> (usize, Option<usize>) {
496        let lim = self.limit - self.index;
497        (0, Some(lim))
498    }
499}
idalib/idb.rs

idalib/
idb.rs
