wraith/manipulation/spoof/
gadget.rs

1//! Gadget finder for legitimate return addresses
2//!
3//! Scans system modules (ntdll, kernel32, etc.) for code gadgets that can be
4//! used as legitimate-looking return addresses for syscall spoofing.
5//!
6//! # Builder API
7//!
8//! The recommended way to find gadgets is using the builder pattern:
9//!
10//! ```ignore
11//! // find jmp rbx gadgets
12//! let gadgets = GadgetFinder::new()?
13//!     .jmp(Register::Rbx)
14//!     .in_module("ntdll.dll")
15//!     .find()?;
16//!
17//! // find any jmp with wildcard
18//! let gadgets = GadgetFinder::new()?
19//!     .pattern("jmp ???")
20//!     .in_module("kernel32.dll")
21//!     .find()?;
22//!
23//! // find specific pattern
24//! let gadgets = GadgetFinder::new()?
25//!     .pattern("jmp rbx")
26//!     .find_in_system_modules()?;
27//! ```
28
29use crate::error::{Result, WraithError};
30use crate::navigation::ModuleQuery;
31use crate::structures::Peb;
32
33#[cfg(feature = "std")]
34use std::collections::HashMap;
35
36#[cfg(all(not(feature = "std"), feature = "alloc"))]
37use alloc::{collections::BTreeMap, format, string::String, vec::Vec};
38
39#[cfg(feature = "std")]
40use std::{format, string::String, vec::Vec};
41
42#[cfg(feature = "std")]
43use std::sync::OnceLock;
44
45#[cfg(feature = "std")]
46static GADGET_CACHE: OnceLock<Result<GadgetCache>> = OnceLock::new();
47
48/// initialize the global gadget cache
49#[cfg(feature = "std")]
50pub fn init_global_cache() -> Result<()> {
51    let result = GADGET_CACHE.get_or_init(GadgetCache::build);
52    match result {
53        Ok(_) => Ok(()),
54        Err(e) => Err(WraithError::SyscallEnumerationFailed {
55            reason: format!("failed to build gadget cache: {}", e),
56        }),
57    }
58}
59
60/// get global gadget cache reference
61#[cfg(feature = "std")]
62pub fn get_global_cache() -> Result<&'static GadgetCache> {
63    let result = GADGET_CACHE.get_or_init(GadgetCache::build);
64    match result {
65        Ok(cache) => Ok(cache),
66        Err(e) => Err(WraithError::SyscallEnumerationFailed {
67            reason: format!("failed to get gadget cache: {}", e),
68        }),
69    }
70}
71
72/// type of gadget instruction sequence
73#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
74pub enum GadgetType {
75    /// jmp rbx - jump to address in rbx
76    JmpRbx,
77    /// jmp rax - jump to address in rax
78    JmpRax,
79    /// jmp rcx - jump to address in rcx
80    JmpRcx,
81    /// jmp rdx - jump to address in rdx
82    JmpRdx,
83    /// jmp r8 - jump to address in r8
84    JmpR8,
85    /// jmp r9 - jump to address in r9
86    JmpR9,
87    /// jmp [rbx] - indirect jump through rbx
88    JmpIndirectRbx,
89    /// jmp [rax] - indirect jump through rax
90    JmpIndirectRax,
91    /// call rbx - call address in rbx
92    CallRbx,
93    /// call rax - call address in rax
94    CallRax,
95    /// ret - simple return
96    Ret,
97    /// add rsp, N; ret - stack cleanup before return
98    AddRspRet { offset: u8 },
99    /// pop reg; ret - pop register then return
100    PopRet { register: u8 },
101    /// push rbx; ret - push rbx onto stack and return (for setting up returns)
102    PushRbxRet,
103}
104
105impl GadgetType {
106    /// get the bytes that make up this gadget type
107    #[cfg(target_arch = "x86_64")]
108    #[must_use]
109    pub fn bytes(&self) -> &'static [u8] {
110        match self {
111            Self::JmpRbx => &[0xFF, 0xE3],        // jmp rbx
112            Self::JmpRax => &[0xFF, 0xE0],        // jmp rax
113            Self::JmpRcx => &[0xFF, 0xE1],        // jmp rcx
114            Self::JmpRdx => &[0xFF, 0xE2],        // jmp rdx
115            Self::JmpR8 => &[0x41, 0xFF, 0xE0],   // jmp r8
116            Self::JmpR9 => &[0x41, 0xFF, 0xE1],   // jmp r9
117            Self::JmpIndirectRbx => &[0xFF, 0x23], // jmp [rbx]
118            Self::JmpIndirectRax => &[0xFF, 0x20], // jmp [rax]
119            Self::CallRbx => &[0xFF, 0xD3],       // call rbx
120            Self::CallRax => &[0xFF, 0xD0],       // call rax
121            Self::Ret => &[0xC3],                 // ret
122            Self::AddRspRet { .. } => &[],        // variable, handled separately
123            Self::PopRet { .. } => &[],           // variable, handled separately
124            Self::PushRbxRet => &[0x53, 0xC3],    // push rbx; ret
125        }
126    }
127
128    /// get friendly name for this gadget type
129    #[must_use]
130    pub fn name(&self) -> &'static str {
131        match self {
132            Self::JmpRbx => "jmp rbx",
133            Self::JmpRax => "jmp rax",
134            Self::JmpRcx => "jmp rcx",
135            Self::JmpRdx => "jmp rdx",
136            Self::JmpR8 => "jmp r8",
137            Self::JmpR9 => "jmp r9",
138            Self::JmpIndirectRbx => "jmp [rbx]",
139            Self::JmpIndirectRax => "jmp [rax]",
140            Self::CallRbx => "call rbx",
141            Self::CallRax => "call rax",
142            Self::Ret => "ret",
143            Self::AddRspRet { offset: _ } => "add rsp, N; ret",
144            Self::PopRet { .. } => "pop reg; ret",
145            Self::PushRbxRet => "push rbx; ret",
146        }
147    }
148
149    /// parse a gadget type from a pattern string
150    /// supports wildcards: "jmp ???" matches any jmp reg
151    #[must_use]
152    pub fn from_pattern(pattern: &str) -> Option<GadgetPattern> {
153        let pattern = pattern.trim().to_lowercase();
154        let parts: Vec<&str> = pattern.split_whitespace().collect();
155
156        if parts.is_empty() {
157            return None;
158        }
159
160        match parts[0] {
161            "jmp" => {
162                if parts.len() < 2 {
163                    return None;
164                }
165                let operand = parts[1];
166                if operand == "???" {
167                    Some(GadgetPattern::JmpAny)
168                } else if let Some(reg) = Register::from_str(operand) {
169                    if operand.starts_with('[') && operand.ends_with(']') {
170                        Some(GadgetPattern::JmpIndirect(reg))
171                    } else {
172                        Some(GadgetPattern::Jmp(reg))
173                    }
174                } else if operand.starts_with('[') && operand.ends_with(']') {
175                    let inner = &operand[1..operand.len()-1];
176                    if inner == "???" {
177                        Some(GadgetPattern::JmpIndirectAny)
178                    } else if let Some(reg) = Register::from_str(inner) {
179                        Some(GadgetPattern::JmpIndirect(reg))
180                    } else {
181                        None
182                    }
183                } else {
184                    None
185                }
186            }
187            "call" => {
188                if parts.len() < 2 {
189                    return None;
190                }
191                let operand = parts[1];
192                if operand == "???" {
193                    Some(GadgetPattern::CallAny)
194                } else if let Some(reg) = Register::from_str(operand) {
195                    Some(GadgetPattern::Call(reg))
196                } else {
197                    None
198                }
199            }
200            "ret" => Some(GadgetPattern::Ret),
201            "pop" => {
202                if parts.len() < 2 {
203                    return Some(GadgetPattern::PopRetAny);
204                }
205                let operand = parts[1].trim_end_matches(';');
206                if operand == "???" {
207                    Some(GadgetPattern::PopRetAny)
208                } else if let Some(reg) = Register::from_str(operand) {
209                    Some(GadgetPattern::PopRet(reg))
210                } else {
211                    None
212                }
213            }
214            "add" => {
215                // "add rsp, ???; ret" or "add rsp, N; ret"
216                if parts.len() >= 3 && parts[1].trim_end_matches(',') == "rsp" {
217                    Some(GadgetPattern::AddRspRet)
218                } else {
219                    None
220                }
221            }
222            "???" => {
223                // wildcard instruction - match the operand
224                if parts.len() < 2 {
225                    Some(GadgetPattern::Any)
226                } else {
227                    let operand = parts[1];
228                    if let Some(reg) = Register::from_str(operand) {
229                        Some(GadgetPattern::AnyWithReg(reg))
230                    } else {
231                        Some(GadgetPattern::Any)
232                    }
233                }
234            }
235            _ => None,
236        }
237    }
238}
239
240/// x86-64 registers for gadget operands
241#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
242pub enum Register {
243    Rax,
244    Rbx,
245    Rcx,
246    Rdx,
247    Rsi,
248    Rdi,
249    Rbp,
250    Rsp,
251    R8,
252    R9,
253    R10,
254    R11,
255    R12,
256    R13,
257    R14,
258    R15,
259}
260
261impl Register {
262    /// parse register from string
263    #[must_use]
264    pub fn from_str(s: &str) -> Option<Self> {
265        match s.to_lowercase().as_str() {
266            "rax" => Some(Self::Rax),
267            "rbx" => Some(Self::Rbx),
268            "rcx" => Some(Self::Rcx),
269            "rdx" => Some(Self::Rdx),
270            "rsi" => Some(Self::Rsi),
271            "rdi" => Some(Self::Rdi),
272            "rbp" => Some(Self::Rbp),
273            "rsp" => Some(Self::Rsp),
274            "r8" => Some(Self::R8),
275            "r9" => Some(Self::R9),
276            "r10" => Some(Self::R10),
277            "r11" => Some(Self::R11),
278            "r12" => Some(Self::R12),
279            "r13" => Some(Self::R13),
280            "r14" => Some(Self::R14),
281            "r15" => Some(Self::R15),
282            _ => None,
283        }
284    }
285
286    /// get the register name
287    #[must_use]
288    pub fn name(&self) -> &'static str {
289        match self {
290            Self::Rax => "rax",
291            Self::Rbx => "rbx",
292            Self::Rcx => "rcx",
293            Self::Rdx => "rdx",
294            Self::Rsi => "rsi",
295            Self::Rdi => "rdi",
296            Self::Rbp => "rbp",
297            Self::Rsp => "rsp",
298            Self::R8 => "r8",
299            Self::R9 => "r9",
300            Self::R10 => "r10",
301            Self::R11 => "r11",
302            Self::R12 => "r12",
303            Self::R13 => "r13",
304            Self::R14 => "r14",
305            Self::R15 => "r15",
306        }
307    }
308}
309
310/// pattern for gadget matching (supports wildcards)
311#[derive(Debug, Clone, Copy, PartialEq, Eq)]
312pub enum GadgetPattern {
313    /// jmp reg
314    Jmp(Register),
315    /// jmp [reg]
316    JmpIndirect(Register),
317    /// any jmp (wildcard)
318    JmpAny,
319    /// any jmp [reg] (wildcard)
320    JmpIndirectAny,
321    /// call reg
322    Call(Register),
323    /// any call (wildcard)
324    CallAny,
325    /// ret
326    Ret,
327    /// add rsp, N; ret
328    AddRspRet,
329    /// pop reg; ret
330    PopRet(Register),
331    /// any pop; ret
332    PopRetAny,
333    /// match any gadget
334    Any,
335    /// match any gadget using specific register
336    AnyWithReg(Register),
337}
338
339impl GadgetPattern {
340    /// check if a gadget type matches this pattern
341    #[must_use]
342    pub fn matches(&self, gadget: &GadgetType) -> bool {
343        match self {
344            Self::Jmp(reg) => match (reg, gadget) {
345                (Register::Rax, GadgetType::JmpRax) => true,
346                (Register::Rbx, GadgetType::JmpRbx) => true,
347                (Register::Rcx, GadgetType::JmpRcx) => true,
348                (Register::Rdx, GadgetType::JmpRdx) => true,
349                (Register::R8, GadgetType::JmpR8) => true,
350                (Register::R9, GadgetType::JmpR9) => true,
351                _ => false,
352            },
353            Self::JmpIndirect(reg) => match (reg, gadget) {
354                (Register::Rax, GadgetType::JmpIndirectRax) => true,
355                (Register::Rbx, GadgetType::JmpIndirectRbx) => true,
356                _ => false,
357            },
358            Self::JmpAny => matches!(
359                gadget,
360                GadgetType::JmpRax
361                    | GadgetType::JmpRbx
362                    | GadgetType::JmpRcx
363                    | GadgetType::JmpRdx
364                    | GadgetType::JmpR8
365                    | GadgetType::JmpR9
366            ),
367            Self::JmpIndirectAny => matches!(
368                gadget,
369                GadgetType::JmpIndirectRax | GadgetType::JmpIndirectRbx
370            ),
371            Self::Call(reg) => match (reg, gadget) {
372                (Register::Rax, GadgetType::CallRax) => true,
373                (Register::Rbx, GadgetType::CallRbx) => true,
374                _ => false,
375            },
376            Self::CallAny => matches!(gadget, GadgetType::CallRax | GadgetType::CallRbx),
377            Self::Ret => matches!(gadget, GadgetType::Ret),
378            Self::AddRspRet => matches!(gadget, GadgetType::AddRspRet { .. }),
379            Self::PopRet(reg) => {
380                if let GadgetType::PopRet { register } = gadget {
381                    Self::reg_index(*reg) == Some(*register)
382                } else {
383                    false
384                }
385            }
386            Self::PopRetAny => matches!(gadget, GadgetType::PopRet { .. }),
387            Self::Any => true,
388            Self::AnyWithReg(reg) => Self::gadget_uses_reg(gadget, *reg),
389        }
390    }
391
392    fn reg_index(reg: Register) -> Option<u8> {
393        match reg {
394            Register::Rax => Some(0),
395            Register::Rcx => Some(1),
396            Register::Rdx => Some(2),
397            Register::Rbx => Some(3),
398            Register::Rsp => Some(4),
399            Register::Rbp => Some(5),
400            Register::Rsi => Some(6),
401            Register::Rdi => Some(7),
402            Register::R8 => Some(8),
403            Register::R9 => Some(9),
404            Register::R10 => Some(10),
405            Register::R11 => Some(11),
406            Register::R12 => Some(12),
407            Register::R13 => Some(13),
408            Register::R14 => Some(14),
409            Register::R15 => Some(15),
410        }
411    }
412
413    fn gadget_uses_reg(gadget: &GadgetType, reg: Register) -> bool {
414        match (gadget, reg) {
415            (GadgetType::JmpRax | GadgetType::JmpIndirectRax | GadgetType::CallRax, Register::Rax) => true,
416            (GadgetType::JmpRbx | GadgetType::JmpIndirectRbx | GadgetType::CallRbx | GadgetType::PushRbxRet, Register::Rbx) => true,
417            (GadgetType::JmpRcx, Register::Rcx) => true,
418            (GadgetType::JmpRdx, Register::Rdx) => true,
419            (GadgetType::JmpR8, Register::R8) => true,
420            (GadgetType::JmpR9, Register::R9) => true,
421            (GadgetType::PopRet { register }, _) => Self::reg_index(reg) == Some(*register),
422            _ => false,
423        }
424    }
425
426    /// get all concrete gadget types matching this pattern
427    #[must_use]
428    pub fn matching_types(&self) -> Vec<GadgetType> {
429        match self {
430            Self::Jmp(reg) => match reg {
431                Register::Rax => vec![GadgetType::JmpRax],
432                Register::Rbx => vec![GadgetType::JmpRbx],
433                Register::Rcx => vec![GadgetType::JmpRcx],
434                Register::Rdx => vec![GadgetType::JmpRdx],
435                Register::R8 => vec![GadgetType::JmpR8],
436                Register::R9 => vec![GadgetType::JmpR9],
437                _ => vec![],
438            },
439            Self::JmpIndirect(reg) => match reg {
440                Register::Rax => vec![GadgetType::JmpIndirectRax],
441                Register::Rbx => vec![GadgetType::JmpIndirectRbx],
442                _ => vec![],
443            },
444            Self::JmpAny => vec![
445                GadgetType::JmpRax,
446                GadgetType::JmpRbx,
447                GadgetType::JmpRcx,
448                GadgetType::JmpRdx,
449                GadgetType::JmpR8,
450                GadgetType::JmpR9,
451            ],
452            Self::JmpIndirectAny => vec![GadgetType::JmpIndirectRax, GadgetType::JmpIndirectRbx],
453            Self::Call(reg) => match reg {
454                Register::Rax => vec![GadgetType::CallRax],
455                Register::Rbx => vec![GadgetType::CallRbx],
456                _ => vec![],
457            },
458            Self::CallAny => vec![GadgetType::CallRax, GadgetType::CallRbx],
459            Self::Ret => vec![GadgetType::Ret],
460            Self::AddRspRet => vec![], // handled specially (variable bytes)
461            Self::PopRet(_) => vec![], // handled specially (variable bytes)
462            Self::PopRetAny => vec![], // handled specially (variable bytes)
463            Self::Any => vec![
464                GadgetType::JmpRax,
465                GadgetType::JmpRbx,
466                GadgetType::JmpRcx,
467                GadgetType::JmpRdx,
468                GadgetType::JmpR8,
469                GadgetType::JmpR9,
470                GadgetType::JmpIndirectRax,
471                GadgetType::JmpIndirectRbx,
472                GadgetType::CallRax,
473                GadgetType::CallRbx,
474                GadgetType::Ret,
475                GadgetType::PushRbxRet,
476            ],
477            Self::AnyWithReg(reg) => {
478                let mut types = vec![];
479                for t in [
480                    GadgetType::JmpRax,
481                    GadgetType::JmpRbx,
482                    GadgetType::JmpRcx,
483                    GadgetType::JmpRdx,
484                    GadgetType::JmpR8,
485                    GadgetType::JmpR9,
486                    GadgetType::JmpIndirectRax,
487                    GadgetType::JmpIndirectRbx,
488                    GadgetType::CallRax,
489                    GadgetType::CallRbx,
490                    GadgetType::PushRbxRet,
491                ] {
492                    if Self::gadget_uses_reg(&t, *reg) {
493                        types.push(t);
494                    }
495                }
496                types
497            }
498        }
499    }
500}
501
502/// a found gadget with its location and type
503#[derive(Debug, Clone)]
504pub struct Gadget {
505    /// absolute address of the gadget
506    pub address: usize,
507    /// type of gadget
508    pub gadget_type: GadgetType,
509    /// module containing this gadget
510    pub module_name: String,
511    /// offset within the module
512    pub module_offset: usize,
513    /// is this in a system module (more trustworthy)
514    pub is_system_module: bool,
515}
516
517impl Gadget {
518    /// check if gadget is still valid (bytes haven't changed)
519    #[must_use]
520    pub fn is_valid(&self) -> bool {
521        let bytes = self.gadget_type.bytes();
522        if bytes.is_empty() {
523            return true; // variable-length gadgets need special handling
524        }
525
526        // SAFETY: we're reading from a previously-validated code address
527        let actual = unsafe { core::slice::from_raw_parts(self.address as *const u8, bytes.len()) };
528        actual == bytes
529    }
530}
531
532/// jmp-type gadget (for jumping to syscall stub)
533#[derive(Debug, Clone)]
534pub struct JmpGadget {
535    pub gadget: Gadget,
536}
537
538impl JmpGadget {
539    #[must_use]
540    pub fn address(&self) -> usize {
541        self.gadget.address
542    }
543}
544
545/// ret-type gadget (for return address spoofing)
546#[derive(Debug, Clone)]
547pub struct RetGadget {
548    pub gadget: Gadget,
549    /// number of bytes the ret pops (for add rsp, N; ret patterns)
550    pub stack_adjustment: usize,
551}
552
553impl RetGadget {
554    #[must_use]
555    pub fn address(&self) -> usize {
556        self.gadget.address
557    }
558}
559
560/// cache of found gadgets organized by type and module
561#[cfg(feature = "std")]
562#[derive(Debug)]
563pub struct GadgetCache {
564    /// gadgets indexed by type
565    by_type: HashMap<GadgetType, Vec<Gadget>>,
566    /// gadgets indexed by module name (lowercase)
567    by_module: HashMap<String, Vec<Gadget>>,
568    /// preferred jmp rbx gadget in ntdll
569    preferred_jmp_rbx: Option<Gadget>,
570    /// preferred jmp rax gadget in ntdll
571    preferred_jmp_rax: Option<Gadget>,
572    /// preferred ret gadget in kernel32
573    preferred_ret: Option<Gadget>,
574}
575
576#[cfg(feature = "std")]
577impl GadgetCache {
578    /// build gadget cache by scanning system modules
579    pub fn build() -> Result<Self> {
580        let finder = GadgetFinder::new()?;
581
582        let mut by_type: HashMap<GadgetType, Vec<Gadget>> = HashMap::new();
583        let mut by_module: HashMap<String, Vec<Gadget>> = HashMap::new();
584
585        // scan key system modules
586        let modules = ["ntdll.dll", "kernel32.dll", "kernelbase.dll"];
587
588        for module_name in modules {
589            if let Ok(gadgets) = finder.scan_module_all(module_name) {
590                for gadget in gadgets {
591                    let module_lower = gadget.module_name.to_lowercase();
592
593                    by_type
594                        .entry(gadget.gadget_type)
595                        .or_default()
596                        .push(gadget.clone());
597
598                    by_module.entry(module_lower).or_default().push(gadget);
599                }
600            }
601        }
602
603        // find preferred gadgets
604        let preferred_jmp_rbx = by_type
605            .get(&GadgetType::JmpRbx)
606            .and_then(|v| v.iter().find(|g| g.module_name.eq_ignore_ascii_case("ntdll.dll")))
607            .cloned();
608
609        let preferred_jmp_rax = by_type
610            .get(&GadgetType::JmpRax)
611            .and_then(|v| v.iter().find(|g| g.module_name.eq_ignore_ascii_case("ntdll.dll")))
612            .cloned();
613
614        let preferred_ret = by_type
615            .get(&GadgetType::Ret)
616            .and_then(|v| {
617                v.iter()
618                    .find(|g| g.module_name.eq_ignore_ascii_case("kernel32.dll"))
619            })
620            .cloned();
621
622        Ok(Self {
623            by_type,
624            by_module,
625            preferred_jmp_rbx,
626            preferred_jmp_rax,
627            preferred_ret,
628        })
629    }
630
631    /// get preferred jmp rbx gadget (in ntdll)
632    #[must_use]
633    pub fn jmp_rbx(&self) -> Option<&Gadget> {
634        self.preferred_jmp_rbx.as_ref()
635    }
636
637    /// get preferred jmp rax gadget (in ntdll)
638    #[must_use]
639    pub fn jmp_rax(&self) -> Option<&Gadget> {
640        self.preferred_jmp_rax.as_ref()
641    }
642
643    /// get preferred ret gadget (in kernel32)
644    #[must_use]
645    pub fn ret_gadget(&self) -> Option<&Gadget> {
646        self.preferred_ret.as_ref()
647    }
648
649    /// get all gadgets of a specific type
650    #[must_use]
651    pub fn get_by_type(&self, gadget_type: GadgetType) -> &[Gadget] {
652        self.by_type.get(&gadget_type).map(|v| v.as_slice()).unwrap_or(&[])
653    }
654
655    /// get all gadgets in a specific module
656    #[must_use]
657    pub fn get_by_module(&self, module_name: &str) -> &[Gadget] {
658        self.by_module
659            .get(&module_name.to_lowercase())
660            .map(|v| v.as_slice())
661            .unwrap_or(&[])
662    }
663
664    /// get first available jmp gadget (tries rbx, then rax)
665    #[must_use]
666    pub fn any_jmp_gadget(&self) -> Option<&Gadget> {
667        self.preferred_jmp_rbx
668            .as_ref()
669            .or(self.preferred_jmp_rax.as_ref())
670            .or_else(|| {
671                self.by_type
672                    .get(&GadgetType::JmpRbx)
673                    .and_then(|v| v.first())
674            })
675            .or_else(|| {
676                self.by_type
677                    .get(&GadgetType::JmpRax)
678                    .and_then(|v| v.first())
679            })
680    }
681
682    /// search gadgets using a pattern
683    #[must_use]
684    pub fn find_by_pattern(&self, pattern: &GadgetPattern) -> Vec<&Gadget> {
685        let mut results = Vec::new();
686        for (gadget_type, gadgets) in &self.by_type {
687            if pattern.matches(gadget_type) {
688                results.extend(gadgets.iter());
689            }
690        }
691        results
692    }
693}
694
695/// scanner for finding gadgets in loaded modules
696pub struct GadgetFinder {
697    peb: Peb,
698}
699
700impl GadgetFinder {
701    /// create new gadget finder
702    pub fn new() -> Result<Self> {
703        Ok(Self {
704            peb: Peb::current()?,
705        })
706    }
707
708    /// find all jmp rbx gadgets in a module
709    pub fn find_jmp_rbx(&self, module_name: &str) -> Result<Vec<JmpGadget>> {
710        self.find_gadgets_of_type(module_name, GadgetType::JmpRbx)
711            .map(|gadgets| gadgets.into_iter().map(|g| JmpGadget { gadget: g }).collect())
712    }
713
714    /// find all jmp rax gadgets in a module
715    pub fn find_jmp_rax(&self, module_name: &str) -> Result<Vec<JmpGadget>> {
716        self.find_gadgets_of_type(module_name, GadgetType::JmpRax)
717            .map(|gadgets| gadgets.into_iter().map(|g| JmpGadget { gadget: g }).collect())
718    }
719
720    /// find all ret gadgets in a module
721    pub fn find_ret(&self, module_name: &str) -> Result<Vec<RetGadget>> {
722        self.find_gadgets_of_type(module_name, GadgetType::Ret)
723            .map(|gadgets| {
724                gadgets
725                    .into_iter()
726                    .map(|g| RetGadget {
727                        gadget: g,
728                        stack_adjustment: 0,
729                    })
730                    .collect()
731            })
732    }
733
734    /// find gadgets of a specific type in a module
735    pub fn find_gadgets_of_type(
736        &self,
737        module_name: &str,
738        gadget_type: GadgetType,
739    ) -> Result<Vec<Gadget>> {
740        let query = ModuleQuery::new(&self.peb);
741        let module = query.find_by_name(module_name)?;
742
743        let bytes = gadget_type.bytes();
744        if bytes.is_empty() {
745            return Ok(Vec::new());
746        }
747
748        let base = module.base();
749        let size = module.size();
750        let name = module.name();
751        let is_system = is_system_module(&name);
752
753        // scan for the byte pattern
754        // SAFETY: module memory is mapped and readable
755        let data = unsafe { core::slice::from_raw_parts(base as *const u8, size) };
756
757        let mut gadgets = Vec::new();
758        let pattern_len = bytes.len();
759
760        // scan for gadget bytes
761        for offset in 0..=(size.saturating_sub(pattern_len)) {
762            if &data[offset..offset + pattern_len] == bytes {
763                gadgets.push(Gadget {
764                    address: base + offset,
765                    gadget_type,
766                    module_name: name.clone(),
767                    module_offset: offset,
768                    is_system_module: is_system,
769                });
770            }
771        }
772
773        Ok(gadgets)
774    }
775
776    /// find add rsp, N; ret gadgets (for stack cleanup)
777    pub fn find_add_rsp_ret(&self, module_name: &str) -> Result<Vec<RetGadget>> {
778        let query = ModuleQuery::new(&self.peb);
779        let module = query.find_by_name(module_name)?;
780
781        let base = module.base();
782        let size = module.size();
783        let name = module.name();
784        let is_system = is_system_module(&name);
785
786        // SAFETY: module memory is mapped and readable
787        let data = unsafe { core::slice::from_raw_parts(base as *const u8, size) };
788
789        let mut gadgets = Vec::new();
790
791        // patterns for add rsp, imm8; ret
792        // 48 83 C4 XX C3 = add rsp, XX; ret (5 bytes)
793        for offset in 0..=(size.saturating_sub(5)) {
794            if data[offset] == 0x48
795                && data[offset + 1] == 0x83
796                && data[offset + 2] == 0xC4
797                && data[offset + 4] == 0xC3
798            {
799                let stack_adj = data[offset + 3] as usize;
800                gadgets.push(RetGadget {
801                    gadget: Gadget {
802                        address: base + offset,
803                        gadget_type: GadgetType::AddRspRet {
804                            offset: data[offset + 3],
805                        },
806                        module_name: name.clone(),
807                        module_offset: offset,
808                        is_system_module: is_system,
809                    },
810                    stack_adjustment: stack_adj,
811                });
812            }
813        }
814
815        // also look for add rsp, imm32; ret
816        // 48 81 C4 XX XX XX XX C3 = add rsp, XXXXXXXX; ret (8 bytes)
817        for offset in 0..=(size.saturating_sub(8)) {
818            if data[offset] == 0x48
819                && data[offset + 1] == 0x81
820                && data[offset + 2] == 0xC4
821                && data[offset + 7] == 0xC3
822            {
823                let stack_adj = u32::from_le_bytes([
824                    data[offset + 3],
825                    data[offset + 4],
826                    data[offset + 5],
827                    data[offset + 6],
828                ]) as usize;
829
830                gadgets.push(RetGadget {
831                    gadget: Gadget {
832                        address: base + offset,
833                        gadget_type: GadgetType::AddRspRet {
834                            offset: 0, // too large for u8
835                        },
836                        module_name: name.clone(),
837                        module_offset: offset,
838                        is_system_module: is_system,
839                    },
840                    stack_adjustment: stack_adj,
841                });
842            }
843        }
844
845        Ok(gadgets)
846    }
847
848    /// find pop reg; ret gadgets
849    pub fn find_pop_ret(&self, module_name: &str) -> Result<Vec<RetGadget>> {
850        let query = ModuleQuery::new(&self.peb);
851        let module = query.find_by_name(module_name)?;
852
853        let base = module.base();
854        let size = module.size();
855        let name = module.name();
856        let is_system = is_system_module(&name);
857
858        // SAFETY: module memory is mapped and readable
859        let data = unsafe { core::slice::from_raw_parts(base as *const u8, size) };
860
861        let mut gadgets = Vec::new();
862
863        // pop rax; ret = 58 C3
864        // pop rcx; ret = 59 C3
865        // pop rdx; ret = 5A C3
866        // pop rbx; ret = 5B C3
867        // pop rsp; ret = 5C C3 (dangerous, skip)
868        // pop rbp; ret = 5D C3
869        // pop rsi; ret = 5E C3
870        // pop rdi; ret = 5F C3
871        for offset in 0..=(size.saturating_sub(2)) {
872            let first = data[offset];
873            if (0x58..=0x5F).contains(&first) && first != 0x5C && data[offset + 1] == 0xC3 {
874                gadgets.push(RetGadget {
875                    gadget: Gadget {
876                        address: base + offset,
877                        gadget_type: GadgetType::PopRet {
878                            register: first - 0x58,
879                        },
880                        module_name: name.clone(),
881                        module_offset: offset,
882                        is_system_module: is_system,
883                    },
884                    stack_adjustment: 8, // one pop
885                });
886            }
887        }
888
889        // also check for REX.W pop; ret (pop r8-r15)
890        // 41 58 C3 = pop r8; ret
891        // 41 59 C3 = pop r9; ret
892        // etc.
893        for offset in 0..=(size.saturating_sub(3)) {
894            if data[offset] == 0x41
895                && (0x58..=0x5F).contains(&data[offset + 1])
896                && data[offset + 1] != 0x5C
897                && data[offset + 2] == 0xC3
898            {
899                gadgets.push(RetGadget {
900                    gadget: Gadget {
901                        address: base + offset,
902                        gadget_type: GadgetType::PopRet {
903                            register: data[offset + 1] - 0x58 + 8,
904                        },
905                        module_name: name.clone(),
906                        module_offset: offset,
907                        is_system_module: is_system,
908                    },
909                    stack_adjustment: 8,
910                });
911            }
912        }
913
914        Ok(gadgets)
915    }
916
917    /// scan a module for all gadget types
918    pub fn scan_module_all(&self, module_name: &str) -> Result<Vec<Gadget>> {
919        let mut all_gadgets = Vec::new();
920
921        // basic jmp gadgets
922        for gadget_type in [
923            GadgetType::JmpRbx,
924            GadgetType::JmpRax,
925            GadgetType::JmpRcx,
926            GadgetType::JmpRdx,
927            GadgetType::CallRbx,
928            GadgetType::CallRax,
929            GadgetType::Ret,
930            GadgetType::PushRbxRet,
931        ] {
932            if let Ok(gadgets) = self.find_gadgets_of_type(module_name, gadget_type) {
933                all_gadgets.extend(gadgets);
934            }
935        }
936
937        // add rsp, N; ret gadgets
938        if let Ok(ret_gadgets) = self.find_add_rsp_ret(module_name) {
939            all_gadgets.extend(ret_gadgets.into_iter().map(|r| r.gadget));
940        }
941
942        // pop; ret gadgets
943        if let Ok(pop_gadgets) = self.find_pop_ret(module_name) {
944            all_gadgets.extend(pop_gadgets.into_iter().map(|r| r.gadget));
945        }
946
947        Ok(all_gadgets)
948    }
949
950    /// find the best jmp gadget for syscall spoofing
951    /// prefers ntdll > kernelbase > kernel32
952    pub fn find_best_jmp_gadget(&self) -> Result<JmpGadget> {
953        // try ntdll first (most legitimate for syscalls)
954        if let Ok(gadgets) = self.find_jmp_rbx("ntdll.dll") {
955            if let Some(g) = gadgets.into_iter().next() {
956                return Ok(g);
957            }
958        }
959
960        if let Ok(gadgets) = self.find_jmp_rax("ntdll.dll") {
961            if let Some(g) = gadgets.into_iter().next() {
962                return Ok(g);
963            }
964        }
965
966        // try kernelbase
967        if let Ok(gadgets) = self.find_jmp_rbx("kernelbase.dll") {
968            if let Some(g) = gadgets.into_iter().next() {
969                return Ok(g);
970            }
971        }
972
973        // try kernel32
974        if let Ok(gadgets) = self.find_jmp_rbx("kernel32.dll") {
975            if let Some(g) = gadgets.into_iter().next() {
976                return Ok(g);
977            }
978        }
979
980        Err(WraithError::SyscallEnumerationFailed {
981            reason: "no suitable jmp gadget found".into(),
982        })
983    }
984
985    /// find a ret gadget that looks legitimate
986    pub fn find_best_ret_gadget(&self) -> Result<RetGadget> {
987        // prefer kernel32 ret gadgets (look like normal API returns)
988        if let Ok(gadgets) = self.find_ret("kernel32.dll") {
989            if let Some(g) = gadgets.into_iter().next() {
990                return Ok(g);
991            }
992        }
993
994        if let Ok(gadgets) = self.find_ret("kernelbase.dll") {
995            if let Some(g) = gadgets.into_iter().next() {
996                return Ok(g);
997            }
998        }
999
1000        if let Ok(gadgets) = self.find_ret("ntdll.dll") {
1001            if let Some(g) = gadgets.into_iter().next() {
1002                return Ok(g);
1003            }
1004        }
1005
1006        Err(WraithError::SyscallEnumerationFailed {
1007            reason: "no suitable ret gadget found".into(),
1008        })
1009    }
1010
1011    // ========== Builder API ==========
1012
1013    /// start building a jmp gadget search
1014    ///
1015    /// # Example
1016    /// ```ignore
1017    /// let gadgets = GadgetFinder::new()?
1018    ///     .jmp(Register::Rbx)
1019    ///     .in_module("ntdll.dll")
1020    ///     .find()?;
1021    /// ```
1022    #[must_use]
1023    pub fn jmp(self, register: Register) -> GadgetSearch {
1024        GadgetSearch {
1025            finder: self,
1026            pattern: GadgetPattern::Jmp(register),
1027            module: None,
1028            system_modules_only: false,
1029        }
1030    }
1031
1032    /// start building a jmp [reg] (indirect) gadget search
1033    #[must_use]
1034    pub fn jmp_indirect(self, register: Register) -> GadgetSearch {
1035        GadgetSearch {
1036            finder: self,
1037            pattern: GadgetPattern::JmpIndirect(register),
1038            module: None,
1039            system_modules_only: false,
1040        }
1041    }
1042
1043    /// start building a call gadget search
1044    #[must_use]
1045    pub fn call(self, register: Register) -> GadgetSearch {
1046        GadgetSearch {
1047            finder: self,
1048            pattern: GadgetPattern::Call(register),
1049            module: None,
1050            system_modules_only: false,
1051        }
1052    }
1053
1054    /// start building a ret gadget search
1055    #[must_use]
1056    pub fn ret(self) -> GadgetSearch {
1057        GadgetSearch {
1058            finder: self,
1059            pattern: GadgetPattern::Ret,
1060            module: None,
1061            system_modules_only: false,
1062        }
1063    }
1064
1065    /// start building a pop; ret gadget search
1066    #[must_use]
1067    pub fn pop_ret(self, register: Register) -> GadgetSearch {
1068        GadgetSearch {
1069            finder: self,
1070            pattern: GadgetPattern::PopRet(register),
1071            module: None,
1072            system_modules_only: false,
1073        }
1074    }
1075
1076    /// start building a search using a pattern string
1077    ///
1078    /// # Supported patterns
1079    /// - `"jmp rbx"` - specific jmp
1080    /// - `"jmp ???"` - any jmp reg (wildcard)
1081    /// - `"jmp [rax]"` - indirect jmp
1082    /// - `"call ???"` - any call
1083    /// - `"ret"` - simple return
1084    /// - `"pop ???; ret"` - any pop then ret
1085    /// - `"??? rbx"` - any instruction using rbx
1086    ///
1087    /// # Example
1088    /// ```ignore
1089    /// let gadgets = GadgetFinder::new()?
1090    ///     .pattern("jmp ???")
1091    ///     .in_module("ntdll.dll")
1092    ///     .find()?;
1093    /// ```
1094    pub fn pattern(self, pattern_str: &str) -> Result<GadgetSearch> {
1095        let pattern = GadgetType::from_pattern(pattern_str).ok_or_else(|| {
1096            WraithError::PatternParseFailed {
1097                reason: format!("invalid gadget pattern: {}", pattern_str),
1098            }
1099        })?;
1100
1101        Ok(GadgetSearch {
1102            finder: self,
1103            pattern,
1104            module: None,
1105            system_modules_only: false,
1106        })
1107    }
1108
1109    /// find gadgets matching a pattern in a module
1110    pub fn find_by_pattern(
1111        &self,
1112        module_name: &str,
1113        pattern: &GadgetPattern,
1114    ) -> Result<Vec<Gadget>> {
1115        let mut results = Vec::new();
1116
1117        // handle concrete gadget types
1118        for gadget_type in pattern.matching_types() {
1119            if let Ok(gadgets) = self.find_gadgets_of_type(module_name, gadget_type) {
1120                results.extend(gadgets);
1121            }
1122        }
1123
1124        // handle variable-length patterns specially
1125        match pattern {
1126            GadgetPattern::AddRspRet => {
1127                if let Ok(ret_gadgets) = self.find_add_rsp_ret(module_name) {
1128                    results.extend(ret_gadgets.into_iter().map(|r| r.gadget));
1129                }
1130            }
1131            GadgetPattern::PopRet(reg) => {
1132                if let Ok(pop_gadgets) = self.find_pop_ret(module_name) {
1133                    let reg_idx = GadgetPattern::reg_index(*reg);
1134                    results.extend(
1135                        pop_gadgets
1136                            .into_iter()
1137                            .filter(|g| {
1138                                if let GadgetType::PopRet { register } = g.gadget.gadget_type {
1139                                    reg_idx == Some(register)
1140                                } else {
1141                                    false
1142                                }
1143                            })
1144                            .map(|r| r.gadget),
1145                    );
1146                }
1147            }
1148            GadgetPattern::PopRetAny => {
1149                if let Ok(pop_gadgets) = self.find_pop_ret(module_name) {
1150                    results.extend(pop_gadgets.into_iter().map(|r| r.gadget));
1151                }
1152            }
1153            _ => {}
1154        }
1155
1156        Ok(results)
1157    }
1158}
1159
1160/// builder for gadget searches with fluent API
1161pub struct GadgetSearch {
1162    finder: GadgetFinder,
1163    pattern: GadgetPattern,
1164    module: Option<String>,
1165    system_modules_only: bool,
1166}
1167
1168impl GadgetSearch {
1169    /// search only in a specific module
1170    #[must_use]
1171    pub fn in_module(mut self, module_name: &str) -> Self {
1172        self.module = Some(module_name.to_string());
1173        self
1174    }
1175
1176    /// search only in system modules (ntdll, kernel32, etc.)
1177    #[must_use]
1178    pub fn system_modules_only(mut self) -> Self {
1179        self.system_modules_only = true;
1180        self
1181    }
1182
1183    /// execute the search and return matching gadgets
1184    pub fn find(self) -> Result<Vec<Gadget>> {
1185        if let Some(module_name) = &self.module {
1186            self.finder.find_by_pattern(module_name, &self.pattern)
1187        } else {
1188            self.find_in_system_modules()
1189        }
1190    }
1191
1192    /// search in all common system modules
1193    pub fn find_in_system_modules(self) -> Result<Vec<Gadget>> {
1194        let modules = ["ntdll.dll", "kernel32.dll", "kernelbase.dll"];
1195        let mut all_gadgets = Vec::new();
1196
1197        for module_name in modules {
1198            if let Ok(gadgets) = self.finder.find_by_pattern(module_name, &self.pattern) {
1199                all_gadgets.extend(gadgets);
1200            }
1201        }
1202
1203        if all_gadgets.is_empty() {
1204            Err(WraithError::GadgetNotFound {
1205                gadget_type: "matching pattern",
1206            })
1207        } else {
1208            Ok(all_gadgets)
1209        }
1210    }
1211
1212    /// get the first matching gadget (convenience method)
1213    pub fn find_first(self) -> Result<Gadget> {
1214        self.find()?
1215            .into_iter()
1216            .next()
1217            .ok_or(WraithError::GadgetNotFound {
1218                gadget_type: "matching pattern",
1219            })
1220    }
1221
1222    /// get the underlying pattern
1223    #[must_use]
1224    pub fn get_pattern(&self) -> &GadgetPattern {
1225        &self.pattern
1226    }
1227}
1228
1229/// check if a module is a system module
1230fn is_system_module(name: &str) -> bool {
1231    let lower = name.to_lowercase();
1232    lower == "ntdll.dll"
1233        || lower == "kernel32.dll"
1234        || lower == "kernelbase.dll"
1235        || lower == "user32.dll"
1236        || lower == "gdi32.dll"
1237        || lower == "advapi32.dll"
1238        || lower == "msvcrt.dll"
1239        || lower == "ws2_32.dll"
1240        || lower == "ole32.dll"
1241        || lower == "combase.dll"
1242}
1243
1244#[cfg(test)]
1245mod tests {
1246    use super::*;
1247
1248    #[test]
1249    fn test_find_jmp_rbx_ntdll() {
1250        let finder = GadgetFinder::new().expect("should create finder");
1251        let gadgets = finder.find_jmp_rbx("ntdll.dll").expect("should find gadgets");
1252
1253        // ntdll should have jmp rbx gadgets
1254        assert!(!gadgets.is_empty(), "should find jmp rbx gadgets in ntdll");
1255
1256        // verify first gadget is valid
1257        let first = &gadgets[0];
1258        assert!(first.gadget.is_valid(), "gadget should be valid");
1259        assert!(first.gadget.is_system_module, "should be system module");
1260    }
1261
1262    #[test]
1263    fn test_find_ret_gadgets() {
1264        let finder = GadgetFinder::new().expect("should create finder");
1265        let gadgets = finder.find_ret("kernel32.dll").expect("should find gadgets");
1266
1267        // kernel32 should have many ret gadgets
1268        assert!(!gadgets.is_empty(), "should find ret gadgets in kernel32");
1269
1270        // verify gadget is valid
1271        let first = &gadgets[0];
1272        assert!(first.gadget.is_valid(), "gadget should be valid");
1273    }
1274
1275    #[test]
1276    fn test_find_add_rsp_ret() {
1277        let finder = GadgetFinder::new().expect("should create finder");
1278
1279        if let Ok(gadgets) = finder.find_add_rsp_ret("ntdll.dll") {
1280            // just check we can find them without crashing
1281            for g in gadgets.iter().take(5) {
1282                assert!(g.stack_adjustment > 0, "should have stack adjustment");
1283            }
1284        }
1285    }
1286
1287    #[test]
1288    fn test_gadget_cache() {
1289        let cache = GadgetCache::build().expect("should build cache");
1290
1291        // should have found some gadgets
1292        assert!(cache.jmp_rbx().is_some() || cache.jmp_rax().is_some());
1293    }
1294
1295    #[test]
1296    fn test_best_jmp_gadget() {
1297        let finder = GadgetFinder::new().expect("should create finder");
1298        let gadget = finder.find_best_jmp_gadget().expect("should find gadget");
1299
1300        assert!(gadget.gadget.is_valid(), "best gadget should be valid");
1301    }
1302}