1use 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#[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#[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#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
74pub enum GadgetType {
75 JmpRbx,
77 JmpRax,
79 JmpRcx,
81 JmpRdx,
83 JmpR8,
85 JmpR9,
87 JmpIndirectRbx,
89 JmpIndirectRax,
91 CallRbx,
93 CallRax,
95 Ret,
97 AddRspRet { offset: u8 },
99 PopRet { register: u8 },
101 PushRbxRet,
103}
104
105impl GadgetType {
106 #[cfg(target_arch = "x86_64")]
108 #[must_use]
109 pub fn bytes(&self) -> &'static [u8] {
110 match self {
111 Self::JmpRbx => &[0xFF, 0xE3], Self::JmpRax => &[0xFF, 0xE0], Self::JmpRcx => &[0xFF, 0xE1], Self::JmpRdx => &[0xFF, 0xE2], Self::JmpR8 => &[0x41, 0xFF, 0xE0], Self::JmpR9 => &[0x41, 0xFF, 0xE1], Self::JmpIndirectRbx => &[0xFF, 0x23], Self::JmpIndirectRax => &[0xFF, 0x20], Self::CallRbx => &[0xFF, 0xD3], Self::CallRax => &[0xFF, 0xD0], Self::Ret => &[0xC3], Self::AddRspRet { .. } => &[], Self::PopRet { .. } => &[], Self::PushRbxRet => &[0x53, 0xC3], }
126 }
127
128 #[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 #[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 if parts.len() >= 3 && parts[1].trim_end_matches(',') == "rsp" {
217 Some(GadgetPattern::AddRspRet)
218 } else {
219 None
220 }
221 }
222 "???" => {
223 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#[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 #[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 #[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#[derive(Debug, Clone, Copy, PartialEq, Eq)]
312pub enum GadgetPattern {
313 Jmp(Register),
315 JmpIndirect(Register),
317 JmpAny,
319 JmpIndirectAny,
321 Call(Register),
323 CallAny,
325 Ret,
327 AddRspRet,
329 PopRet(Register),
331 PopRetAny,
333 Any,
335 AnyWithReg(Register),
337}
338
339impl GadgetPattern {
340 #[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 #[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![], Self::PopRet(_) => vec![], Self::PopRetAny => vec![], 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#[derive(Debug, Clone)]
504pub struct Gadget {
505 pub address: usize,
507 pub gadget_type: GadgetType,
509 pub module_name: String,
511 pub module_offset: usize,
513 pub is_system_module: bool,
515}
516
517impl Gadget {
518 #[must_use]
520 pub fn is_valid(&self) -> bool {
521 let bytes = self.gadget_type.bytes();
522 if bytes.is_empty() {
523 return true; }
525
526 let actual = unsafe { core::slice::from_raw_parts(self.address as *const u8, bytes.len()) };
528 actual == bytes
529 }
530}
531
532#[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#[derive(Debug, Clone)]
547pub struct RetGadget {
548 pub gadget: Gadget,
549 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#[cfg(feature = "std")]
562#[derive(Debug)]
563pub struct GadgetCache {
564 by_type: HashMap<GadgetType, Vec<Gadget>>,
566 by_module: HashMap<String, Vec<Gadget>>,
568 preferred_jmp_rbx: Option<Gadget>,
570 preferred_jmp_rax: Option<Gadget>,
572 preferred_ret: Option<Gadget>,
574}
575
576#[cfg(feature = "std")]
577impl GadgetCache {
578 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 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 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 #[must_use]
633 pub fn jmp_rbx(&self) -> Option<&Gadget> {
634 self.preferred_jmp_rbx.as_ref()
635 }
636
637 #[must_use]
639 pub fn jmp_rax(&self) -> Option<&Gadget> {
640 self.preferred_jmp_rax.as_ref()
641 }
642
643 #[must_use]
645 pub fn ret_gadget(&self) -> Option<&Gadget> {
646 self.preferred_ret.as_ref()
647 }
648
649 #[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 #[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 #[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 #[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
695pub struct GadgetFinder {
697 peb: Peb,
698}
699
700impl GadgetFinder {
701 pub fn new() -> Result<Self> {
703 Ok(Self {
704 peb: Peb::current()?,
705 })
706 }
707
708 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 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 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 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 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 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 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 let data = unsafe { core::slice::from_raw_parts(base as *const u8, size) };
788
789 let mut gadgets = Vec::new();
790
791 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 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, },
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 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 let data = unsafe { core::slice::from_raw_parts(base as *const u8, size) };
860
861 let mut gadgets = Vec::new();
862
863 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, });
886 }
887 }
888
889 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 pub fn scan_module_all(&self, module_name: &str) -> Result<Vec<Gadget>> {
919 let mut all_gadgets = Vec::new();
920
921 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 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 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 pub fn find_best_jmp_gadget(&self) -> Result<JmpGadget> {
953 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 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 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 pub fn find_best_ret_gadget(&self) -> Result<RetGadget> {
987 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 #[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 #[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 #[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 #[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 #[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 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 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 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 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
1160pub struct GadgetSearch {
1162 finder: GadgetFinder,
1163 pattern: GadgetPattern,
1164 module: Option<String>,
1165 system_modules_only: bool,
1166}
1167
1168impl GadgetSearch {
1169 #[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 #[must_use]
1178 pub fn system_modules_only(mut self) -> Self {
1179 self.system_modules_only = true;
1180 self
1181 }
1182
1183 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 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 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 #[must_use]
1224 pub fn get_pattern(&self) -> &GadgetPattern {
1225 &self.pattern
1226 }
1227}
1228
1229fn 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 assert!(!gadgets.is_empty(), "should find jmp rbx gadgets in ntdll");
1255
1256 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 assert!(!gadgets.is_empty(), "should find ret gadgets in kernel32");
1269
1270 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 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 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}