memflow_native/windows/
mod.rs

1use memflow::os::process::*;
2use memflow::prelude::v1::*;
3
4use windows::core::PCSTR;
5use windows::Win32::Foundation::{CloseHandle, HANDLE};
6
7use windows::Win32::System::Diagnostics::ToolHelp::{
8    CreateToolhelp32Snapshot, Process32FirstW, Process32NextW, PROCESSENTRY32W, TH32CS_SNAPPROCESS,
9};
10
11use windows::Win32::System::Threading::{GetCurrentProcess, OpenProcessToken};
12
13use windows::Win32::Security::{
14    AdjustTokenPrivileges, LookupPrivilegeValueA, LUID_AND_ATTRIBUTES, SE_PRIVILEGE_ENABLED,
15    TOKEN_ADJUST_PRIVILEGES, TOKEN_PRIVILEGES,
16};
17
18use core::mem::{size_of, MaybeUninit};
19use core::ptr;
20use std::ffi::OsString;
21use std::os::windows::ffi::OsStringExt;
22
23pub mod mem;
24use mem::ProcessVirtualMemory;
25
26pub mod process;
27pub use process::WindowsProcess;
28
29pub mod keyboard;
30pub use keyboard::{WindowsKeyboard, WindowsKeyboardState};
31
32struct KernelModule {}
33
34pub(crate) struct Handle(HANDLE);
35
36impl From<HANDLE> for Handle {
37    fn from(handle: HANDLE) -> Handle {
38        Handle(handle)
39    }
40}
41
42impl core::ops::Deref for Handle {
43    type Target = HANDLE;
44
45    fn deref(&self) -> &Self::Target {
46        &self.0
47    }
48}
49
50impl Drop for Handle {
51    fn drop(&mut self) {
52        unsafe { CloseHandle(self.0) }.ok();
53    }
54}
55
56pub fn conv_err(_err: windows::core::Error) -> Error {
57    // TODO: proper error kind
58    // TODO: proper origin
59    Error(ErrorOrigin::OsLayer, ErrorKind::Unknown)
60}
61
62unsafe fn enable_debug_privilege() -> Result<()> {
63    let process = GetCurrentProcess();
64    let mut token = HANDLE::default();
65
66    OpenProcessToken(process, TOKEN_ADJUST_PRIVILEGES, &mut token).map_err(conv_err)?;
67
68    let mut luid = Default::default();
69
70    let mut se_debug_name = *b"SeDebugPrivilege\0";
71
72    LookupPrivilegeValueA(
73        PCSTR(core::ptr::null_mut()),
74        PCSTR(se_debug_name.as_mut_ptr()),
75        &mut luid,
76    )
77    .map_err(conv_err)?;
78
79    let new_privileges = TOKEN_PRIVILEGES {
80        PrivilegeCount: 1,
81        Privileges: [LUID_AND_ATTRIBUTES {
82            Luid: luid,
83            Attributes: SE_PRIVILEGE_ENABLED,
84        }],
85    };
86
87    AdjustTokenPrivileges(
88        token,
89        false,
90        Some(&new_privileges),
91        std::mem::size_of_val(&new_privileges) as _,
92        None,
93        None,
94    )
95    .map_err(conv_err)
96}
97
98pub struct WindowsOs {
99    info: OsInfo,
100    cached_processes: Vec<ProcessInfo>,
101    cached_modules: Vec<KernelModule>,
102}
103
104impl WindowsOs {
105    pub fn new(args: &OsArgs) -> Result<Self> {
106        match args.extra_args.get("elevate_token") {
107            Some("off") | Some("OFF") | Some("Off") | Some("n") | Some("N") | Some("0") => {}
108            _ => {
109                unsafe { enable_debug_privilege() }?;
110            }
111        }
112
113        Ok(Default::default())
114    }
115}
116
117impl Clone for WindowsOs {
118    fn clone(&self) -> Self {
119        Self {
120            info: self.info.clone(),
121            cached_processes: vec![],
122            cached_modules: vec![],
123        }
124    }
125}
126
127impl Default for WindowsOs {
128    fn default() -> Self {
129        let info = OsInfo {
130            base: Address::NULL,
131            size: 0,
132            arch: ArchitectureIdent::X86(64, false),
133        };
134
135        Self {
136            info,
137            cached_modules: vec![],
138            cached_processes: vec![],
139        }
140    }
141}
142
143impl Os for WindowsOs {
144    type ProcessType<'a> = WindowsProcess;
145    type IntoProcessType = WindowsProcess;
146
147    /// Walks a process list and calls a callback for each process structure address
148    ///
149    /// The callback is fully opaque. We need this style so that C FFI can work seamlessly.
150    fn process_address_list_callback(&mut self, callback: AddressCallback) -> Result<()> {
151        let handle =
152            Handle(unsafe { CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0) }.map_err(conv_err)?);
153
154        let mut maybe_entry = MaybeUninit::<PROCESSENTRY32W>::uninit();
155
156        unsafe {
157            ptr::write(
158                &mut (*maybe_entry.as_mut_ptr()).dwSize,
159                size_of::<PROCESSENTRY32W>() as u32,
160            );
161        }
162
163        let ptr = maybe_entry.as_mut_ptr();
164
165        std::iter::once(unsafe { Process32FirstW(*handle, ptr) })
166            .chain(std::iter::repeat_with(|| unsafe {
167                Process32NextW(*handle, ptr)
168            }))
169            .take_while(|b| b.is_ok())
170            .map(|_| unsafe { maybe_entry.assume_init() })
171            .map(|p| {
172                let address = Address::from(p.th32ProcessID as umem);
173                let len = p.szExeFile.iter().take_while(|&&c| c != 0).count();
174
175                let path = OsString::from_wide(&p.szExeFile[..len]);
176                let path = path.to_string_lossy();
177                let path = &*path;
178                let name = path.rsplit_once('\\').map(|(_, end)| end).unwrap_or(path);
179
180                self.cached_processes.push(ProcessInfo {
181                    address,
182                    pid: address.to_umem() as _,
183                    state: ProcessState::Alive,
184                    name: name.into(),
185                    path: path.into(),
186                    command_line: "".into(),
187                    sys_arch: self.info.arch,
188                    proc_arch: self.info.arch,
189                    // dtb is not known/used here
190                    dtb1: Address::invalid(),
191                    dtb2: Address::invalid(),
192                });
193
194                address
195            })
196            .feed_into(callback);
197
198        Ok(())
199    }
200
201    /// Find process information by its internal address
202    fn process_info_by_address(&mut self, address: Address) -> Result<ProcessInfo> {
203        self.cached_processes
204            .iter()
205            .find(|p| p.address == address)
206            .cloned()
207            .ok_or(Error(ErrorOrigin::OsLayer, ErrorKind::NotFound))
208    }
209
210    /// Construct a process by its info, borrowing the OS
211    ///
212    /// It will share the underlying memory resources
213    fn process_by_info(&mut self, info: ProcessInfo) -> Result<Self::ProcessType<'_>> {
214        WindowsProcess::try_new(info)
215    }
216
217    /// Construct a process by its info, consuming the OS
218    ///
219    /// This function will consume the Kernel instance and move its resources into the process
220    fn into_process_by_info(mut self, info: ProcessInfo) -> Result<Self::IntoProcessType> {
221        self.process_by_info(info)
222    }
223
224    /// Walks the OS module list and calls the provided callback for each module structure
225    /// address
226    ///
227    /// # Arguments
228    /// * `callback` - where to pass each matching module to. This is an opaque callback.
229    fn module_address_list_callback(&mut self, mut callback: AddressCallback) -> Result<()> {
230        (0..self.cached_modules.len())
231            .map(Address::from)
232            .take_while(|a| callback.call(*a))
233            .for_each(|_| {});
234
235        Ok(())
236    }
237
238    /// Retrieves a module by its structure address
239    ///
240    /// # Arguments
241    /// * `address` - address where module's information resides in
242    fn module_by_address(&mut self, _address: Address) -> Result<ModuleInfo> {
243        /*self.cached_modules
244        .iter()
245        .skip(address.to_umem() as usize)
246        .next()
247        .map(|km| ModuleInfo {
248            address,
249            size: km.size as umem,
250            base: Address::NULL,
251            name: km
252                .name
253                .split("/")
254                .last()
255                .or(Some(""))
256                .map(ReprCString::from)
257                .unwrap(),
258            arch: self.info.arch,
259            path: km.name.clone().into(),
260            parent_process: Address::INVALID,
261        })
262        .ok_or(Error(ErrorOrigin::OsLayer, ErrorKind::NotFound))*/
263
264        todo!()
265    }
266
267    /// Retrieves address of the primary module structure of the process
268    ///
269    /// This will generally be for the initial executable that was run
270    fn primary_module_address(&mut self) -> Result<Address> {
271        Ok(self.module_by_name("ntoskrnl.exe")?.address)
272    }
273
274    /// Retrieves information for the primary module of the process
275    ///
276    /// This will generally be the initial executable that was run
277    fn primary_module(&mut self) -> Result<ModuleInfo> {
278        self.module_by_name("ntoskrnl.exe")
279    }
280
281    /// Retrieves a list of all imports of a given module
282    fn module_import_list_callback(
283        &mut self,
284        _info: &ModuleInfo,
285        _callback: ImportCallback,
286    ) -> Result<()> {
287        //memflow::os::util::module_import_list_callback(&mut self.virt_mem, info, callback)
288        Err(Error(ErrorOrigin::OsLayer, ErrorKind::NotImplemented))
289    }
290
291    /// Retrieves a list of all exports of a given module
292    fn module_export_list_callback(
293        &mut self,
294        _info: &ModuleInfo,
295        _callback: ExportCallback,
296    ) -> Result<()> {
297        //memflow::os::util::module_export_list_callback(&mut self.virt_mem, info, callback)
298        Err(Error(ErrorOrigin::OsLayer, ErrorKind::NotImplemented))
299    }
300
301    /// Retrieves a list of all sections of a given module
302    fn module_section_list_callback(
303        &mut self,
304        _info: &ModuleInfo,
305        _callback: SectionCallback,
306    ) -> Result<()> {
307        //memflow::os::util::module_section_list_callback(&mut self.virt_mem, info, callback)
308        Err(Error(ErrorOrigin::OsLayer, ErrorKind::NotImplemented))
309    }
310
311    /// Retrieves the OS info
312    fn info(&self) -> &OsInfo {
313        &self.info
314    }
315}
316
317impl OsKeyboard for WindowsOs {
318    type KeyboardType<'a> = WindowsKeyboard;
319    type IntoKeyboardType = WindowsKeyboard;
320
321    fn keyboard(&mut self) -> memflow::error::Result<Self::KeyboardType<'_>> {
322        Ok(WindowsKeyboard::new())
323    }
324
325    fn into_keyboard(self) -> memflow::error::Result<Self::IntoKeyboardType> {
326        Ok(WindowsKeyboard::new())
327    }
328}