wraith/manipulation/remote/
process.rs

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

wraith/manipulation/remote/
process.rs
