Skip to main content

idalib/
idb.rs

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