wraith/manipulation/remote/
process.rs

1//! Remote process wrapper for cross-process operations
2
3use 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/// process access rights configuration
14#[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
53/// wrapper for a remote process handle with memory operations
54pub struct RemoteProcess {
55    handle: usize,
56    pid: u32,
57    owns_handle: bool,
58}
59
60impl RemoteProcess {
61    /// open a process by PID with specified access rights
62    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    /// open a process with all access rights
81    pub fn open_all_access(pid: u32) -> Result<Self> {
82        Self::open(pid, ProcessAccess::all())
83    }
84
85    /// create from an existing handle (does not take ownership)
86    ///
87    /// # Safety
88    /// caller must ensure handle is valid and has appropriate access rights
89    pub unsafe fn from_handle(handle: usize, pid: u32) -> Self {
90        Self {
91            handle,
92            pid,
93            owns_handle: false,
94        }
95    }
96
97    /// create from an existing handle (takes ownership)
98    ///
99    /// # Safety
100    /// caller must ensure handle is valid and has appropriate access rights
101    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    /// get the raw process handle
110    pub fn handle(&self) -> usize {
111        self.handle
112    }
113
114    /// get the process ID
115    pub fn pid(&self) -> u32 {
116        self.pid
117    }
118
119    /// read memory from the remote process
120    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    /// read a typed value from the remote process
130    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        // SAFETY: buffer is correctly sized and we just read the bytes
134        Ok(unsafe { (buffer.as_ptr() as *const T).read_unaligned() })
135    }
136
137    /// read a null-terminated string from the remote process
138    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    /// read a wide string from the remote process
152    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    /// write memory to the remote process
173    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    /// write a typed value to the remote process
183    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    /// allocate memory in the remote process
194    pub fn allocate(
195        &self,
196        size: usize,
197        protection: u32,
198    ) -> Result<RemoteAllocation> {
199        self.allocate_at(0, size, protection)
200    }
201
202    /// allocate memory at a preferred address
203    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    /// allocate RW memory
230    pub fn allocate_rw(&self, size: usize) -> Result<RemoteAllocation> {
231        self.allocate(size, PAGE_READWRITE)
232    }
233
234    /// allocate RWX memory
235    pub fn allocate_rwx(&self, size: usize) -> Result<RemoteAllocation> {
236        self.allocate(size, PAGE_EXECUTE_READWRITE)
237    }
238
239    /// allocate RX memory
240    pub fn allocate_rx(&self, size: usize) -> Result<RemoteAllocation> {
241        self.allocate(size, PAGE_EXECUTE_READ)
242    }
243
244    /// change memory protection in the remote process
245    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    /// change protection with RAII guard that restores on drop
255    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    /// free allocated memory in the remote process
271    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    /// write shellcode and allocate executable memory
281    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    /// write and execute shellcode via remote thread
289    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
309// SAFETY: handle is process-specific and can be sent between threads
310unsafe impl Send for RemoteProcess {}
311unsafe impl Sync for RemoteProcess {}
312
313/// RAII wrapper for remote memory allocation
314pub struct RemoteAllocation {
315    process_handle: usize,
316    base: usize,
317    size: usize,
318    owns_memory: bool,
319}
320
321impl RemoteAllocation {
322    /// get the base address
323    pub fn base(&self) -> usize {
324        self.base
325    }
326
327    /// get the allocation size
328    pub fn size(&self) -> usize {
329        self.size
330    }
331
332    /// leak the allocation (don't free on drop)
333    pub fn leak(mut self) -> usize {
334        self.owns_memory = false;
335        self.base
336    }
337
338    /// create from raw values without ownership
339    ///
340    /// # Safety
341    /// caller must ensure the memory region is valid
342    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
360/// RAII guard for remote memory protection changes
361pub struct RemoteProtectionGuard {
362    process_handle: usize,
363    address: usize,
364    size: usize,
365    old_protection: u32,
366}
367
368impl RemoteProtectionGuard {
369    /// get the old protection that will be restored
370    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}