wraith/manipulation/remote/
process.rs1use crate::error::{Result, WraithError};
4use crate::manipulation::syscall::{
5 nt_allocate_virtual_memory, nt_close, nt_free_virtual_memory, nt_open_process,
6 nt_protect_virtual_memory, nt_read_virtual_memory, nt_write_virtual_memory,
7 ClientId, ObjectAttributes, MEM_COMMIT, MEM_RELEASE, MEM_RESERVE,
8 PAGE_EXECUTE_READ, PAGE_EXECUTE_READWRITE, PAGE_READWRITE,
9 PROCESS_ALL_ACCESS, PROCESS_QUERY_INFORMATION, PROCESS_VM_OPERATION,
10 PROCESS_VM_READ, PROCESS_VM_WRITE,
11};
12
13#[derive(Debug, Clone, Copy)]
15pub struct ProcessAccess {
16 pub rights: u32,
17}
18
19impl ProcessAccess {
20 pub const fn all() -> Self {
21 Self { rights: PROCESS_ALL_ACCESS }
22 }
23
24 pub const fn read_write() -> Self {
25 Self {
26 rights: PROCESS_VM_READ | PROCESS_VM_WRITE | PROCESS_VM_OPERATION | PROCESS_QUERY_INFORMATION,
27 }
28 }
29
30 pub const fn read_only() -> Self {
31 Self {
32 rights: PROCESS_VM_READ | PROCESS_QUERY_INFORMATION,
33 }
34 }
35
36 pub const fn query() -> Self {
37 Self {
38 rights: PROCESS_QUERY_INFORMATION,
39 }
40 }
41
42 pub const fn custom(rights: u32) -> Self {
43 Self { rights }
44 }
45}
46
47impl Default for ProcessAccess {
48 fn default() -> Self {
49 Self::all()
50 }
51}
52
53pub struct RemoteProcess {
55 handle: usize,
56 pid: u32,
57 owns_handle: bool,
58}
59
60impl RemoteProcess {
61 pub fn open(pid: u32, access: ProcessAccess) -> Result<Self> {
63 let obj_attr = ObjectAttributes::new();
64 let client_id = ClientId::for_process(pid);
65
66 let handle = nt_open_process(access.rights, &obj_attr, &client_id).map_err(|e| {
67 WraithError::ProcessOpenFailed {
68 pid,
69 reason: format!("{}", e),
70 }
71 })?;
72
73 Ok(Self {
74 handle,
75 pid,
76 owns_handle: true,
77 })
78 }
79
80 pub fn open_all_access(pid: u32) -> Result<Self> {
82 Self::open(pid, ProcessAccess::all())
83 }
84
85 pub unsafe fn from_handle(handle: usize, pid: u32) -> Self {
90 Self {
91 handle,
92 pid,
93 owns_handle: false,
94 }
95 }
96
97 pub unsafe fn from_handle_owned(handle: usize, pid: u32) -> Self {
102 Self {
103 handle,
104 pid,
105 owns_handle: true,
106 }
107 }
108
109 pub fn handle(&self) -> usize {
111 self.handle
112 }
113
114 pub fn pid(&self) -> u32 {
116 self.pid
117 }
118
119 pub fn read(&self, address: usize, buffer: &mut [u8]) -> Result<usize> {
121 nt_read_virtual_memory(self.handle, address, buffer).map_err(|e| {
122 WraithError::ReadFailed {
123 address: address as u64,
124 size: buffer.len(),
125 }
126 })
127 }
128
129 pub fn read_value<T: Copy>(&self, address: usize) -> Result<T> {
131 let mut buffer = vec![0u8; core::mem::size_of::<T>()];
132 self.read(address, &mut buffer)?;
133 Ok(unsafe { (buffer.as_ptr() as *const T).read_unaligned() })
135 }
136
137 pub fn read_string(&self, address: usize, max_len: usize) -> Result<String> {
139 let mut buffer = vec![0u8; max_len];
140 let bytes_read = self.read(address, &mut buffer)?;
141
142 let end = buffer.iter()
143 .take(bytes_read)
144 .position(|&b| b == 0)
145 .unwrap_or(bytes_read);
146
147 String::from_utf8_lossy(&buffer[..end]).into_owned();
148 Ok(String::from_utf8_lossy(&buffer[..end]).into_owned())
149 }
150
151 pub fn read_wstring(&self, address: usize, max_chars: usize) -> Result<String> {
153 let mut buffer = vec![0u16; max_chars];
154 let byte_buffer = unsafe {
155 core::slice::from_raw_parts_mut(
156 buffer.as_mut_ptr() as *mut u8,
157 max_chars * 2,
158 )
159 };
160
161 let bytes_read = self.read(address, byte_buffer)?;
162 let chars_read = bytes_read / 2;
163
164 let end = buffer.iter()
165 .take(chars_read)
166 .position(|&c| c == 0)
167 .unwrap_or(chars_read);
168
169 Ok(String::from_utf16_lossy(&buffer[..end]))
170 }
171
172 pub fn write(&self, address: usize, buffer: &[u8]) -> Result<usize> {
174 nt_write_virtual_memory(self.handle, address, buffer).map_err(|_| {
175 WraithError::WriteFailed {
176 address: address as u64,
177 size: buffer.len(),
178 }
179 })
180 }
181
182 pub fn write_value<T: Copy>(&self, address: usize, value: &T) -> Result<usize> {
184 let bytes = unsafe {
185 core::slice::from_raw_parts(
186 value as *const T as *const u8,
187 core::mem::size_of::<T>(),
188 )
189 };
190 self.write(address, bytes)
191 }
192
193 pub fn allocate(
195 &self,
196 size: usize,
197 protection: u32,
198 ) -> Result<RemoteAllocation> {
199 self.allocate_at(0, size, protection)
200 }
201
202 pub fn allocate_at(
204 &self,
205 preferred_base: usize,
206 size: usize,
207 protection: u32,
208 ) -> Result<RemoteAllocation> {
209 let (base, actual_size) = nt_allocate_virtual_memory(
210 self.handle,
211 preferred_base,
212 size,
213 MEM_COMMIT | MEM_RESERVE,
214 protection,
215 )
216 .map_err(|e| WraithError::AllocationFailed {
217 size,
218 protection,
219 })?;
220
221 Ok(RemoteAllocation {
222 process_handle: self.handle,
223 base,
224 size: actual_size,
225 owns_memory: true,
226 })
227 }
228
229 pub fn allocate_rw(&self, size: usize) -> Result<RemoteAllocation> {
231 self.allocate(size, PAGE_READWRITE)
232 }
233
234 pub fn allocate_rwx(&self, size: usize) -> Result<RemoteAllocation> {
236 self.allocate(size, PAGE_EXECUTE_READWRITE)
237 }
238
239 pub fn allocate_rx(&self, size: usize) -> Result<RemoteAllocation> {
241 self.allocate(size, PAGE_EXECUTE_READ)
242 }
243
244 pub fn protect(&self, address: usize, size: usize, protection: u32) -> Result<u32> {
246 nt_protect_virtual_memory(self.handle, address, size, protection).map_err(|_| {
247 WraithError::ProtectionChangeFailed {
248 address: address as u64,
249 size,
250 }
251 })
252 }
253
254 pub fn protect_guard(
256 &self,
257 address: usize,
258 size: usize,
259 new_protection: u32,
260 ) -> Result<RemoteProtectionGuard> {
261 let old_protection = self.protect(address, size, new_protection)?;
262 Ok(RemoteProtectionGuard {
263 process_handle: self.handle,
264 address,
265 size,
266 old_protection,
267 })
268 }
269
270 pub fn free(&self, address: usize) -> Result<()> {
272 nt_free_virtual_memory(self.handle, address, MEM_RELEASE).map_err(|_| {
273 WraithError::AllocationFailed {
274 size: 0,
275 protection: 0,
276 }
277 })
278 }
279
280 pub fn write_shellcode(&self, shellcode: &[u8]) -> Result<RemoteAllocation> {
282 let alloc = self.allocate_rw(shellcode.len())?;
283 self.write(alloc.base, shellcode)?;
284 self.protect(alloc.base, alloc.size, PAGE_EXECUTE_READ)?;
285 Ok(alloc)
286 }
287
288 pub fn execute_shellcode(&self, shellcode: &[u8]) -> Result<u32> {
290 let alloc = self.write_shellcode(shellcode)?;
291 let thread = super::create_remote_thread(
292 self,
293 alloc.base,
294 0,
295 super::RemoteThreadOptions::default(),
296 )?;
297 Ok(thread.id())
298 }
299}
300
301impl Drop for RemoteProcess {
302 fn drop(&mut self) {
303 if self.owns_handle && self.handle != 0 {
304 let _ = nt_close(self.handle);
305 }
306 }
307}
308
309unsafe impl Send for RemoteProcess {}
311unsafe impl Sync for RemoteProcess {}
312
313pub struct RemoteAllocation {
315 process_handle: usize,
316 base: usize,
317 size: usize,
318 owns_memory: bool,
319}
320
321impl RemoteAllocation {
322 pub fn base(&self) -> usize {
324 self.base
325 }
326
327 pub fn size(&self) -> usize {
329 self.size
330 }
331
332 pub fn leak(mut self) -> usize {
334 self.owns_memory = false;
335 self.base
336 }
337
338 pub unsafe fn from_raw(process_handle: usize, base: usize, size: usize) -> Self {
343 Self {
344 process_handle,
345 base,
346 size,
347 owns_memory: false,
348 }
349 }
350}
351
352impl Drop for RemoteAllocation {
353 fn drop(&mut self) {
354 if self.owns_memory && self.base != 0 {
355 let _ = nt_free_virtual_memory(self.process_handle, self.base, MEM_RELEASE);
356 }
357 }
358}
359
360pub struct RemoteProtectionGuard {
362 process_handle: usize,
363 address: usize,
364 size: usize,
365 old_protection: u32,
366}
367
368impl RemoteProtectionGuard {
369 pub fn old_protection(&self) -> u32 {
371 self.old_protection
372 }
373}
374
375impl Drop for RemoteProtectionGuard {
376 fn drop(&mut self) {
377 let _ = nt_protect_virtual_memory(
378 self.process_handle,
379 self.address,
380 self.size,
381 self.old_protection,
382 );
383 }
384}
385
386#[cfg(test)]
387mod tests {
388 use super::*;
389
390 #[test]
391 fn test_process_access_constants() {
392 let all = ProcessAccess::all();
393 assert_eq!(all.rights, PROCESS_ALL_ACCESS);
394
395 let rw = ProcessAccess::read_write();
396 assert!(rw.rights & PROCESS_VM_READ != 0);
397 assert!(rw.rights & PROCESS_VM_WRITE != 0);
398 }
399
400 #[test]
401 fn test_open_current_process() {
402 let pid = std::process::id();
403 let proc = RemoteProcess::open(pid, ProcessAccess::read_only());
404 assert!(proc.is_ok(), "should open current process");
405
406 let proc = proc.unwrap();
407 assert_eq!(proc.pid(), pid);
408 assert!(proc.handle() != 0);
409 }
410}