1use std::{mem, ptr};
4
5use windows::{
6 Win32::{
7 Foundation::{
8 CloseHandle, ERROR_SUCCESS, GetLastError, HANDLE, LUID, SetLastError, WIN32_ERROR,
9 },
10 Security::{
11 AdjustTokenPrivileges, DuplicateTokenEx, GetSidSubAuthority, GetSidSubAuthorityCount,
12 GetTokenInformation, ImpersonateLoggedOnUser, IsTokenRestricted, LUID_AND_ATTRIBUTES,
13 LookupPrivilegeValueW, RevertToSelf, SE_PRIVILEGE_ENABLED, SecurityImpersonation,
14 SetTokenInformation, TOKEN_ACCESS_MASK, TOKEN_ADJUST_PRIVILEGES, TOKEN_ALL_ACCESS,
15 TOKEN_GROUPS, TOKEN_INFORMATION_CLASS, TOKEN_MANDATORY_LABEL, TOKEN_PRIVILEGES,
16 TOKEN_QUERY, TOKEN_USER, TokenGroups, TokenImpersonation, TokenIntegrityLevel,
17 TokenPrimary, TokenPrivileges, TokenSessionId, TokenUser,
18 },
19 System::{
20 SystemServices::MAXIMUM_ALLOWED,
21 Threading::{
22 GetCurrentProcess, OpenProcess, OpenProcessToken, PROCESS_QUERY_INFORMATION,
23 },
24 },
25 },
26 core::PCWSTR,
27};
28
29use super::{HandleGuard, Sid, privilege::Privilege, win32};
30use crate::error::ElevateError;
31
32#[derive(Debug, Clone, Copy, PartialEq, Eq)]
34pub enum TokenType {
35 Primary,
37 Impersonation,
39}
40
41pub struct Token(HANDLE);
43
44impl Drop for Token {
45 fn drop(&mut self) {
46 if !self.0.is_invalid() {
47 unsafe {
48 let _ = CloseHandle(self.0);
49 }
50 }
51 }
52}
53
54impl Token {
55 pub fn current() -> Result<Self, ElevateError> {
57 let mut handle = HANDLE::default();
58 win32!(OpenProcessToken(
59 GetCurrentProcess(),
60 TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY,
61 &raw mut handle
62 ))?;
63 Ok(Self(handle))
64 }
65
66 pub fn from_pid(pid: u32) -> Result<Self, ElevateError> {
68 let proc = HandleGuard(win32!(OpenProcess(PROCESS_QUERY_INFORMATION, false, pid))?);
69 let mut token = HANDLE::default();
70 win32!(OpenProcessToken(
71 proc.0,
72 TOKEN_ACCESS_MASK(MAXIMUM_ALLOWED),
73 &raw mut token
74 ))?;
75 Ok(Self(token))
76 }
77
78 pub const fn as_raw(&self) -> HANDLE {
80 self.0
81 }
82
83 pub fn duplicate(&self, token_type: TokenType) -> Result<Self, ElevateError> {
85 let access = match token_type {
86 TokenType::Primary => TOKEN_ALL_ACCESS,
87 TokenType::Impersonation => TOKEN_ACCESS_MASK(MAXIMUM_ALLOWED),
88 };
89 let raw_type = match token_type {
90 TokenType::Primary => TokenPrimary,
91 TokenType::Impersonation => TokenImpersonation,
92 };
93
94 let mut dup = HANDLE::default();
95 win32!(DuplicateTokenEx(
96 self.0,
97 access,
98 None,
99 SecurityImpersonation,
100 raw_type,
101 &raw mut dup
102 ))?;
103 Ok(Self(dup))
104 }
105
106 pub fn enable_privileges(&self, privileges: &[Privilege]) -> Result<(), ElevateError> {
108 for &privilege in privileges {
109 let mut luid = LUID::default();
110 win32!(LookupPrivilegeValueW(
111 PCWSTR::null(),
112 privilege.as_pcwstr(),
113 &raw mut luid
114 ))?;
115
116 let privs = TOKEN_PRIVILEGES {
117 PrivilegeCount: 1,
118 Privileges: [LUID_AND_ATTRIBUTES {
119 Luid: luid,
120 Attributes: SE_PRIVILEGE_ENABLED,
121 }],
122 };
123
124 unsafe { SetLastError(WIN32_ERROR(0)) };
125 win32!(AdjustTokenPrivileges(
126 self.0,
127 false,
128 Some(ptr::from_ref(&privs)),
129 mem::size_of::<TOKEN_PRIVILEGES>() as u32,
130 None,
131 None,
132 ))?;
133
134 if unsafe { GetLastError() } != ERROR_SUCCESS {
135 return Err(ElevateError::PrivilegeNotAssigned(privilege.as_str()));
136 }
137 }
138 Ok(())
139 }
140
141 pub fn is_restricted(&self) -> Result<bool, ElevateError> {
146 Ok(unsafe { IsTokenRestricted(self.0) }.is_ok())
147 }
148
149 pub fn integrity_level(&self) -> Result<u32, ElevateError> {
151 let buf = self.query_info_buffer(TokenIntegrityLevel)?;
152 let label = unsafe { &*buf.as_ptr().cast::<TOKEN_MANDATORY_LABEL>() };
153
154 let count_ptr = unsafe { GetSidSubAuthorityCount(label.Label.Sid) };
155 if count_ptr.is_null() {
156 return Err(ElevateError::InvalidTokenBuffer);
157 }
158 let count = unsafe { *count_ptr };
159 if count == 0 {
160 return Err(ElevateError::InvalidTokenBuffer);
161 }
162
163 let sub_auth_ptr = unsafe { GetSidSubAuthority(label.Label.Sid, count as u32 - 1) };
164 if sub_auth_ptr.is_null() {
165 return Err(ElevateError::InvalidTokenBuffer);
166 }
167
168 Ok(unsafe { *sub_auth_ptr })
169 }
170
171 pub fn has_privilege_present(&self, privilege: Privilege) -> Result<bool, ElevateError> {
174 let mut target_luid = LUID::default();
175 win32!(LookupPrivilegeValueW(
176 PCWSTR::null(),
177 privilege.as_pcwstr(),
178 &raw mut target_luid
179 ))?;
180
181 let buf = self.query_info_buffer(TokenPrivileges)?;
182 let privs = unsafe { &*buf.as_ptr().cast::<TOKEN_PRIVILEGES>() };
183 let slice = unsafe {
184 std::slice::from_raw_parts(privs.Privileges.as_ptr(), privs.PrivilegeCount as usize)
185 };
186
187 Ok(slice.iter().any(|p| p.Luid == target_luid))
188 }
189
190 pub fn session_id(&self) -> Result<u32, ElevateError> {
192 let mut session_id = 0u32;
193 let mut len = 0u32;
194 win32!(GetTokenInformation(
195 self.0,
196 TokenSessionId,
197 Some(ptr::from_mut(&mut session_id).cast()),
198 mem::size_of::<u32>() as u32,
199 &raw mut len
200 ))?;
201 Ok(session_id)
202 }
203
204 pub fn set_session_id(&self, session_id: u32) -> Result<(), ElevateError> {
206 win32!(SetTokenInformation(
207 self.0,
208 TokenSessionId,
209 ptr::from_ref(&session_id).cast(),
210 mem::size_of::<u32>() as u32
211 ))?;
212 Ok(())
213 }
214
215 pub fn impersonate(&self) -> Result<ImpersonationGuard, ElevateError> {
217 win32!(ImpersonateLoggedOnUser(self.0))?;
218 Ok(ImpersonationGuard)
219 }
220
221 pub fn user_sid(&self) -> Result<String, ElevateError> {
223 let buf = self.query_info_buffer(TokenUser)?;
224 let user = unsafe { &*buf.as_ptr().cast::<TOKEN_USER>() };
226 Sid::string_from_raw(user.User.Sid)
227 }
228
229 pub fn has_sid(&self, target_sid: PCWSTR) -> Result<bool, ElevateError> {
231 let target = Sid::parse(target_sid)?;
232 let buf = self.query_info_buffer(TokenGroups)?;
233 let groups = unsafe { &*buf.as_ptr().cast::<TOKEN_GROUPS>() };
235
236 let slice = unsafe {
237 std::slice::from_raw_parts(groups.Groups.as_ptr(), groups.GroupCount as usize)
238 };
239 Ok(slice.iter().any(|g| Sid::equal(g.Sid, target.as_raw())))
240 }
241
242 fn query_info_buffer(
244 &self,
245 class: TOKEN_INFORMATION_CLASS,
246 ) -> Result<Vec<usize>, ElevateError> {
247 let mut len = 0u32;
248 let _ = unsafe { GetTokenInformation(self.0, class, None, 0, &raw mut len) };
249 if len == 0 {
250 return Err(ElevateError::InvalidTokenBuffer);
251 }
252
253 let count = (len as usize).div_ceil(mem::size_of::<usize>());
254 let mut buf: Vec<usize> = vec![0; count];
255
256 win32!(GetTokenInformation(
257 self.0,
258 class,
259 Some(buf.as_mut_ptr().cast()),
260 len,
261 &raw mut len
262 ))?;
263 Ok(buf)
264 }
265}
266
267pub struct ImpersonationGuard;
269
270impl Drop for ImpersonationGuard {
271 fn drop(&mut self) {
272 unsafe {
273 let _ = RevertToSelf();
274 }
275 }
276}