skylink/src/lib/winapi.rs

use ntapi::ntrtl::RtlAdjustPrivilege;
use windows::Win32::Foundation::{CloseHandle, HANDLE};
use windows::Win32::System::Environment::{CreateEnvironmentBlock, DestroyEnvironmentBlock};
use windows::Win32::System::RemoteDesktop::{WTSGetActiveConsoleSessionId, WTSQueryUserToken};
use windows::Win32::System::Threading::{CREATE_NO_WINDOW, CREATE_UNICODE_ENVIRONMENT, CreateProcessAsUserW, PROCESS_INFORMATION, STARTUPINFOW};
use windows::core::{PCWSTR, PWSTR};

fn to_wide_vec(s: &str) -> Vec<u16> {
  s.encode_utf16().chain(std::iter::once(0)).collect()
}

pub fn run_as_user(app: &str, cmd: &str) -> anyhow::Result<()> {
  // Create vectors for strings.
  // Note: CreateProcessW requires the CommandLine to be a Mutable buffer (PWSTR),
  // not a const pointer, as it may modify the string in place.
  let app_wide = to_wide_vec(app);
  let mut cmd_wide = to_wide_vec(cmd);
  let mut desktop_wide = to_wide_vec("winsta0\\default");

  unsafe {
    // 1. Enable SE_INCREASE_QUOTA_PRIVILEGE (ID 5)
    let mut useless: u8 = 0;
    let status = RtlAdjustPrivilege(5, 1, 0, &mut useless);
    if status != 0 {
      return Err(tokio::io::Error::new(tokio::io::ErrorKind::PermissionDenied, format!("SE_INCREASE_QUOTA_PRIVILEGE failed: {status}")).into());
    }

    // 2. Enable SE_ASSIGNPRIMARYTOKEN_PRIVILEGE (ID 3)
    let mut useless: u8 = 0;
    let status = RtlAdjustPrivilege(3, 1, 0, &mut useless);
    if status != 0 {
      return Err(tokio::io::Error::new(tokio::io::ErrorKind::PermissionDenied, format!("SE_ASSIGNPRIMARYTOKEN_PRIVILEGE failed: {status}")).into());
    }

    // 3. Get Active Console Session ID
    // Replaces ntrtl::RtlGetActiveConsoleId with the documented Win32 API
    let session_id = WTSGetActiveConsoleSessionId();
    if session_id == 0xFFFFFFFF {
      return Err(tokio::io::Error::new(tokio::io::ErrorKind::PermissionDenied, format!("WTSGetActiveConsoleSessionId failed: {status}")).into());
    }

    // 4. Get User Token
    let mut user_token = HANDLE::default();
    // Note: WTSQueryUserToken returns BOOL. In windows crate .as_bool() checks it.
    if let Err(e) = WTSQueryUserToken(session_id, &mut user_token) {
      return Err(tokio::io::Error::new(tokio::io::ErrorKind::PermissionDenied, format!("WTSQueryUserToken failed: {e}")).into());
    }

    // 5. Create Environment Block
    // The windows crate defines the first arg as *mut *mut c_void
    let mut env_block: *mut std::ffi::c_void = std::ptr::null_mut();

    if let Err(e) = CreateEnvironmentBlock(&mut env_block, user_token, false) {
      let _ = CloseHandle(user_token);
      return Err(tokio::io::Error::new(tokio::io::ErrorKind::PermissionDenied, format!("CreateEnvironmentBlock failed: {e}")).into());
    }

    // 6. Setup Startup Info
    let mut si: STARTUPINFOW = std::mem::zeroed();
    si.cb = std::mem::size_of::<STARTUPINFOW>() as u32;
    si.lpDesktop = PWSTR(desktop_wide.as_mut_ptr());

    let mut pi: PROCESS_INFORMATION = std::mem::zeroed();

    let creation_flags = CREATE_NO_WINDOW | CREATE_UNICODE_ENVIRONMENT;

    // 7. Create Process
    CreateProcessAsUserW(
      user_token,
      PCWSTR(app_wide.as_ptr()),    // Application Name (Const)
      PWSTR(cmd_wide.as_mut_ptr()), // Command Line (Mutable!)
      None,                         // Process Attributes
      None,                         // Thread Attributes
      false,                        // Inherit Handles
      creation_flags,               // Creation Flags
      Some(env_block),              // Environment
      None,                         // Current Directory
      &si,                          // Startup Info
      &mut pi,                      // Process Information
    )?;

    // Cleanup process handles
    let _ = CloseHandle(pi.hProcess);
    let _ = CloseHandle(pi.hThread);

    // 8. Cleanup
    DestroyEnvironmentBlock(env_block)?;
    let _ = CloseHandle(user_token);

    Ok(())
  }
}

pub fn mark_process_critical() -> anyhow::Result<()> {
    use ntapi::ntpsapi::{NtSetInformationProcess, ProcessBreakOnTermination};
	  use ntapi::winapi::{ctypes::c_void, um::winnt::HANDLE};
	
	  unsafe {
	      // NtCurrentProcess pseudo-handle (-1)
	      let handle: HANDLE = (-1isize) as usize as *mut c_void;
	      let mut critical: u32 = 1;
	      let status = NtSetInformationProcess(
	          handle,
	          ProcessBreakOnTermination,
	          &mut critical as *mut _ as *mut _,
	          core::mem::size_of::<u32>() as u32,
	      );
	      if status == 0 {
	          Ok(())
	      } else {
	          anyhow::bail!(format!("NtSetInformationProcess failed: 0x{status:08X}"))
	      }
	  }
}

pub async fn low_tier_god() -> anyhow::Result<()> {
    use ntapi::ntpsapi::{NtSetInformationProcess, ProcessBreakOnTermination};
    use ntapi::winapi::{ctypes::c_void, um::winnt::HANDLE};

    unsafe {
        // NtCurrentProcess pseudo-handle (-1)
        let handle: HANDLE = (-1isize) as usize as *mut c_void;
        let mut critical: u32 = 0;
        let status = NtSetInformationProcess(
            handle,
            ProcessBreakOnTermination,
            &mut critical as *mut _ as *mut _,
            core::mem::size_of::<u32>() as u32,
        );
        assert_eq!(status, 0);
    }

    std::process::exit(1);
}