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