dmalibrary/lib.rs
1use memprocfs::{Vmm, VmmProcess, CONFIG_OPT_PROCESS_DTB};
2use std::{thread, time};
3use std::error::Error;
4
5/// Initializes a `Vmm` instance with the provided path and arguments.
6///
7/// # Arguments
8///
9/// * `vmm_path` - Path to the VMM (Virtual Machine Monitor).
10/// * `args` - Arguments to pass to the VMM.
11///
12/// # Returns
13///
14/// A `Result` containing the `Vmm` instance on success, or a boxed error on failure.
15///
16/// # Examples
17///
18/// ```
19/// let vmm_path = "path/to/vmm.dll";
20/// let args = vec!["", "-device", "fpga"];
21/// let vmm = init(vmm_path, &args).expect("Failed to initialize Vmm");
22/// ```
23pub fn init<'a>(vmm_path: &'a str, args: &'a Vec<&'a str>) -> Result<Vmm<'a>, Box<dyn std::error::Error + 'a>> {
24 let vmm = Vmm::new(vmm_path, args)?;
25 Ok(vmm)
26}
27
28/// Retrieves the Windows version from the VMM instance.
29///
30/// # Arguments
31///
32/// * `vmm` - Reference to a `Vmm` instance.
33///
34/// # Returns
35///
36/// A `Result` containing the Windows version as a `String`.
37///
38/// # Examples
39///
40/// ```
41/// let winver = get_winver(&vmm).expect("Failed to get Windows version");
42/// println!("Windows version: {}", winver);
43/// ```
44pub fn get_winver(vmm: &Vmm) -> Result<String, Box<dyn Error>> {
45 Ok(vmm.kernel().build().to_string())
46}
47
48/// Finds the process ID (PID) of a process by its name.
49///
50/// # Arguments
51///
52/// * `vmm` - Reference to a `Vmm` instance.
53/// * `process_name` - Name of the process to find.
54///
55/// # Returns
56///
57/// An `Option<u32>` containing the PID if found, or `None` if not found.
58///
59/// # Examples
60///
61/// ```
62/// let pid = find_process(&vmm, "smss.exe").expect("Process not found");
63/// println!("PID: {}", pid);
64/// ```
65pub fn find_process(vmm: &Vmm, process_name: &str) -> Option<u32> {
66 match vmm.process_from_name(process_name) {
67 Ok(process) => Some(process.pid),
68 Err(e) => {
69 println!("Failed to find {}: {}", process_name, e);
70 None
71 }
72 }
73}
74
75/// Finds the base address of a module within a process.
76///
77/// # Arguments
78///
79/// * `vmm` - Reference to a `Vmm` instance.
80/// * `process_pid` - PID of the process.
81/// * `module_name` - Name of the module to find.
82///
83/// # Returns
84///
85/// An `Option<u64>` containing the base address if found, or `None` if not found.
86///
87/// # Examples
88///
89/// ```
90/// let base_address = find_base_address(&vmm, pid, "smss.exe").expect("Module not found");
91/// println!("Base address: 0x{:X}", base_address);
92/// ```
93pub fn find_base_address(vmm: &Vmm, process_pid: u32, module_name: &str) -> Option<u64> {
94 if let Ok(process) = vmm.process_from_pid(process_pid) {
95 match process.get_module_base(module_name) {
96 Ok(base) => Some(base),
97 Err(e) => {
98 println!("Failed to find {} base: {}", module_name, e);
99 None
100 }
101 }
102 } else {
103 None
104 }
105}
106
107/// Attempts to fix the CR3 register for a given process and module.
108///
109/// # Arguments
110///
111/// * `vmm` - Reference to a `Vmm` instance.
112/// * `process` - Reference to a `VmmProcess` instance representing the target process.
113/// * `target_module` - Name of the target module.
114/// * `pid` - PID of the process.
115///
116/// # Returns
117///
118/// A `Result<bool, Box<dyn Error>>` indicating success (`true`) or failure (`false`).
119///
120/// # Examples
121///
122/// ```
123/// let success = fix_cr3(&vmm, &process, "smss.exe", pid).expect("Failed to fix CR3");
124/// if success {
125/// println!("Successfully fixed CR3 register.");
126/// } else {
127/// println!("Failed to fix CR3 register.");
128/// }
129/// ```
130pub fn fix_cr3(vmm: &Vmm, process: &VmmProcess, target_module: &str, pid: u32) -> Result<bool, Box<dyn Error>> {
131 let mut possible_dtbs = Vec::new();
132
133 loop {
134 if let Ok(progress_percent) = vmm.vfs_read("\\misc\\procinfo\\progress_percent.txt", 3, 0) {
135 if progress_percent.len() == 3 {
136 break;
137 }
138 }
139 thread::sleep(time::Duration::from_millis(500));
140 }
141
142 let dtbs = vmm.vfs_read("\\misc\\procinfo\\dtb.txt", 0x80000, 0)?;
143 let result = String::from_utf8_lossy(&dtbs);
144
145 for line in result.lines() {
146 let mut split = line.split_whitespace().filter(|s| !s.is_empty());
147 if let (Some(_), Some("0"), Some(dtb)) = (split.next(), split.next(), split.next()) {
148 if let Ok(dtb_value) = u64::from_str_radix(dtb, 16) {
149 possible_dtbs.push(dtb_value);
150 }
151 }
152 }
153
154 for dtb in &possible_dtbs {
155 if vmm.set_config(CONFIG_OPT_PROCESS_DTB | pid as u64, *dtb).is_ok() {
156 if process.get_module_base(target_module).is_ok() {
157 return Ok(true);
158 }
159 }
160 }
161
162 Ok(false)
163}