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

idalib/
idb.rs
