elevate_code/
util.rs

1use std::ffi::CString;
2
3use windows::{
4    core::{PCSTR, PSTR},
5    Win32::{
6        Foundation::CloseHandle,
7        System::Threading::{
8            CreateProcessA, ResumeThread, TerminateProcess, WaitForSingleObject, CREATE_SUSPENDED,
9            INFINITE, PROCESS_INFORMATION, STARTUPINFOA,
10        },
11    },
12};
13
14pub struct CommandLineBuilder {
15    args: Vec<String>,
16}
17
18impl CommandLineBuilder {
19    pub fn new() -> Self {
20        Self { args: Vec::new() }
21    }
22
23    pub fn arg(mut self, arg: &str) -> Self {
24        self.args.push(arg.to_string());
25        self
26    }
27
28    pub fn encode(&self) -> String {
29        let mut params = String::new();
30        for arg in self.args.iter() {
31            params.push(' ');
32            if arg.len() == 0 {
33                params.push_str("\"\"");
34            } else if arg.find(&[' ', '\t', '"'][..]).is_none() {
35                params.push_str(&arg);
36            } else {
37                params.push('"');
38                for c in arg.chars() {
39                    match c {
40                        '\\' => params.push_str("\\\\"),
41                        '"' => params.push_str("\\\""),
42                        c => params.push(c),
43                    }
44                }
45                params.push('"');
46            }
47        }
48        if !params.is_empty() {
49            params.remove(0);
50        }
51        params
52    }
53}
54
55pub enum ProcessControlFlow {
56    ResumeMainThread,
57    Terminate,
58}
59
60pub fn create_process(args: &[&str], work: impl Fn(u32) -> ProcessControlFlow) {
61    unsafe {
62        let mut si = STARTUPINFOA::default();
63        let mut pi = PROCESS_INFORMATION::default();
64        si.cb = std::mem::size_of_val(&si) as _;
65        let exe = std::env::current_exe()
66            .unwrap()
67            .to_str()
68            .unwrap()
69            .to_string();
70        let mut builder = CommandLineBuilder::new().arg(&exe);
71        for i in args {
72            builder = builder.arg(i);
73        }
74        let cmd = CString::new(builder.encode()).unwrap();
75        CreateProcessA(
76            PCSTR::null(),
77            PSTR::from_raw(cmd.as_ptr() as _),
78            None,
79            None,
80            true,
81            CREATE_SUSPENDED,
82            None,
83            PCSTR::null(),
84            &si,
85            &mut pi,
86        )
87        .unwrap();
88
89        match work(pi.dwProcessId) {
90            ProcessControlFlow::ResumeMainThread => {
91                ResumeThread(pi.hThread);
92                WaitForSingleObject(pi.hProcess, INFINITE);
93            }
94            ProcessControlFlow::Terminate => {
95                let _ = TerminateProcess(pi.hProcess, u32::MAX);
96            }
97        }
98
99        CloseHandle(pi.hThread).unwrap();
100        CloseHandle(pi.hProcess).unwrap();
101    };
102}
elevate_code/util.rs

elevate_code/
util.rs
