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