idalib/
idb.rs

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

idalib/
idb.rs
