From 843cc696452f6b2d097037c8cde5e4a633f6d00e Mon Sep 17 00:00:00 2001
From: Jake Shadle <jake.shadle@embark-studios.com>
Date: Mon, 18 Nov 2024 15:21:53 +0100
Subject: [PATCH] Mmk

---
 src/bindings.toml   |  3 ++-
 src/lib.rs          | 61 +++++++++++++++++++++++++++++++++++----------
 src/win_bindings.rs | 33 ++++++++++++++++--------
 3 files changed, 73 insertions(+), 24 deletions(-)

diff --git a/src/bindings.toml b/src/bindings.toml
index e9d457d..236f757 100644
--- a/src/bindings.toml
+++ b/src/bindings.toml
@@ -1,8 +1,9 @@
 output = "win_bindings.rs"
 binds = [
     "MAX_PATH",
+    "NtClose",
+    "NtOpenProcess",
     "NtQueryInformationProcess",
-    "OpenProcess",
     "ProcessBasicInformation",
     "ProcessImageFileName",
     "PROCESS_BASIC_INFORMATION",
diff --git a/src/lib.rs b/src/lib.rs
index 2f78964..4fb15b9 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -271,54 +271,92 @@ pub fn is_powershell_parent() -> bool {
     use std::os::windows::ffi::OsStringExt as _;
     use win_bindings::*;
 
-    let mut pid = Some(-1 /* NtCurrentProcess */);
+    struct NtHandle {
+        handle: isize,
+    }
+
+    impl Drop for NtHandle {
+        fn drop(&mut self) {
+            if self.handle != -1 {
+                unsafe {
+                    nt_close(self.handle);
+                }
+            }
+        }
+    }
+
+    let mut handle = Some(NtHandle { handle: -1 });
 
     unsafe {
         let reset = |fname: &mut [u16]| {
             let ustr = &mut *fname.as_mut_ptr().cast::<UnicodeString>();
             ustr.length = 0;
             ustr.maximum_length = MaxPath as _;
-            ustr.buffer = fname
-                .as_mut_ptr()
-                .byte_offset(std::mem::size_of::<UnicodeString>() as _);
         };
 
         // The API for this is extremely irritating, the struct and string buffer
         // need to be the same :/
         let mut file_name = [0u16; MaxPath as usize + std::mem::size_of::<UnicodeString>() / 2];
 
-        while let Some(ppid) = pid {
+        while let Some(ph) = handle {
             let mut basic_info = std::mem::MaybeUninit::<ProcessBasicInformation>::uninit();
             let mut length = 0;
             if dbg!(nt_query_information_process(
-                ppid,
+                ph.handle,
                 Processinfoclass::ProcessBasicInformation,
                 basic_info.as_mut_ptr().cast(),
                 std::mem::size_of::<ProcessBasicInformation>() as _,
                 &mut length,
             )) != StatusSuccess
             {
-                return false;
+                break;
             }
 
             if length != std::mem::size_of::<ProcessBasicInformation>() as u32 {
-                return false;
+                break;
             }
 
             let basic_info = basic_info.assume_init();
             reset(&mut file_name);
 
+            let ppid = basic_info.inherited_from_unique_process_id as isize;
+
+            if ppid == 0 || ppid == -1 {
+                break;
+            }
+
+            let mut parent_handle = -1;
+            let obj_attr = std::mem::zeroed();
+            let client_id = ClientId {
+                unique_process: ppid,
+                unique_thread: 0,
+            };
+            if dbg!(nt_open_process(
+                &parent_handle,
+                ProcessAccessRights::ProcessQueryInformation,
+                &obj_attr,
+                &client_id
+            )) != StatusSuccess
+            {
+                break;
+            }
+
+            handle = Some(NtHandle {
+                handle: parent_handle,
+            });
+
             if dbg!(nt_query_information_process(
-                basic_info.inherited_from_unique_process_id as _,
+                parent_handle,
                 Processinfoclass::ProcessImageFileName,
                 file_name.as_mut_ptr().cast(),
                 (file_name.len() * 2) as _,
                 &mut length,
             )) != StatusSuccess
             {
-                return false;
+                break;
             }
 
+            dbg!(length);
             let ustr = &*file_name.as_ptr().cast::<UnicodeString>();
             let os = std::ffi::OsString::from_wide(
                 &file_name[std::mem::size_of::<UnicodeString>() * 2
@@ -331,9 +369,6 @@ pub fn is_powershell_parent() -> bool {
             if p.file_stem() == Some(std::ffi::OsStr::new("pwsh")) {
                 return true;
             }
-
-            pid = (basic_info.inherited_from_unique_process_id != 0)
-                .then_some(basic_info.inherited_from_unique_process_id as isize);
         }
 
         false
diff --git a/src/win_bindings.rs b/src/win_bindings.rs
index 78fb2ee..eab6f1d 100644
--- a/src/win_bindings.rs
+++ b/src/win_bindings.rs
@@ -5,17 +5,17 @@
     non_camel_case_types,
     clippy::upper_case_acronyms
 )]
-#[link(name = "kernel32", kind = "raw-dylib")]
-extern "system" {
-    #[link_name = "OpenProcess"]
-    pub fn open_process(
-        desired_access: ProcessAccessRights::Enum,
-        inherit_handle: Bool,
-        process_id: u32,
-    ) -> Handle;
-}
 #[link(name = "ntdll", kind = "raw-dylib")]
 extern "system" {
+    #[link_name = "NtClose"]
+    pub fn nt_close(handle: Handle) -> Ntstatus;
+    #[link_name = "NtOpenProcess"]
+    pub fn nt_open_process(
+        process_handle: *mut Handle,
+        desired_access: u32,
+        object_attributes: *const ObjectAttributes,
+        client_id: *const ClientId,
+    ) -> Ntstatus;
     #[link_name = "NtQueryInformationProcess"]
     pub fn nt_query_information_process(
         process_handle: Handle,
@@ -26,7 +26,11 @@ extern "system" {
     ) -> Ntstatus;
 }
 pub const MaxPath: u32 = 260;
-pub type Bool = i32;
+#[repr(C)]
+pub struct ClientId {
+    pub unique_process: Handle,
+    pub unique_thread: Handle,
+}
 pub type Handle = isize;
 #[repr(C)]
 pub struct ListEntry {
@@ -36,6 +40,15 @@ pub struct ListEntry {
 pub type Ntstatus = i32;
 pub const StatusSuccess: Ntstatus = 0;
 #[repr(C)]
+pub struct ObjectAttributes {
+    pub length: u32,
+    pub root_directory: Handle,
+    pub object_name: *mut UnicodeString,
+    pub attributes: u32,
+    pub security_descriptor: *mut ::core::ffi::c_void,
+    pub security_quality_of_service: *mut ::core::ffi::c_void,
+}
+#[repr(C)]
 pub struct Peb {
     pub reserved1: [u8; 2],
     pub being_debugged: u8,