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