wraith/structures/
teb.rs

1//! TEB (Thread Environment Block) structure
2
3use super::offsets::TebOffsets;
4use crate::arch::segment;
5use crate::error::{Result, WraithError};
6use crate::version::WindowsVersion;
7use core::ptr::NonNull;
8
9/// safe wrapper around TEB access
10pub struct Teb {
11    ptr: NonNull<u8>,
12    offsets: &'static TebOffsets,
13}
14
15impl Teb {
16    /// get TEB for current thread
17    pub fn current() -> Result<Self> {
18        let ptr = unsafe { segment::get_teb() };
19        let ptr = NonNull::new(ptr).ok_or(WraithError::InvalidTebAccess)?;
20
21        let version = WindowsVersion::current()?;
22        let offsets = TebOffsets::for_version(&version)?;
23
24        Ok(Self { ptr, offsets })
25    }
26
27    /// raw TEB pointer
28    pub fn as_ptr(&self) -> *mut u8 {
29        self.ptr.as_ptr()
30    }
31
32    /// get process ID
33    pub fn process_id(&self) -> u32 {
34        unsafe {
35            let addr = self.ptr.as_ptr().add(self.offsets.client_id);
36            // ClientId.UniqueProcess is first field
37            #[cfg(target_arch = "x86_64")]
38            {
39                *(addr as *const u64) as u32
40            }
41            #[cfg(target_arch = "x86")]
42            {
43                *(addr as *const u32)
44            }
45        }
46    }
47
48    /// get thread ID
49    pub fn thread_id(&self) -> u32 {
50        unsafe {
51            let addr = self.ptr.as_ptr().add(self.offsets.client_id);
52            // ClientId.UniqueThread is second field
53            #[cfg(target_arch = "x86_64")]
54            {
55                let tid_addr = addr.add(8);
56                *(tid_addr as *const u64) as u32
57            }
58            #[cfg(target_arch = "x86")]
59            {
60                let tid_addr = addr.add(4);
61                *(tid_addr as *const u32)
62            }
63        }
64    }
65
66    /// get PEB pointer from TEB
67    pub fn peb(&self) -> *mut u8 {
68        unsafe {
69            let addr = self.ptr.as_ptr().add(self.offsets.peb);
70            #[cfg(target_arch = "x86_64")]
71            {
72                *(addr as *const u64) as *mut u8
73            }
74            #[cfg(target_arch = "x86")]
75            {
76                *(addr as *const u32) as *mut u8
77            }
78        }
79    }
80
81    /// get last error value
82    pub fn last_error(&self) -> u32 {
83        unsafe {
84            let addr = self.ptr.as_ptr().add(self.offsets.last_error);
85            *(addr as *const u32)
86        }
87    }
88
89    /// set last error value
90    pub fn set_last_error(&mut self, value: u32) {
91        unsafe {
92            let addr = self.ptr.as_ptr().add(self.offsets.last_error);
93            *(addr as *mut u32) = value;
94        }
95    }
96
97    /// get stack base
98    pub fn stack_base(&self) -> usize {
99        unsafe {
100            let addr = self.ptr.as_ptr().add(self.offsets.stack_base);
101            #[cfg(target_arch = "x86_64")]
102            {
103                *(addr as *const u64) as usize
104            }
105            #[cfg(target_arch = "x86")]
106            {
107                *(addr as *const u32) as usize
108            }
109        }
110    }
111
112    /// get stack limit
113    pub fn stack_limit(&self) -> usize {
114        unsafe {
115            let addr = self.ptr.as_ptr().add(self.offsets.stack_limit);
116            #[cfg(target_arch = "x86_64")]
117            {
118                *(addr as *const u64) as usize
119            }
120            #[cfg(target_arch = "x86")]
121            {
122                *(addr as *const u32) as usize
123            }
124        }
125    }
126}