1 /**
2  * Windows process utility code.
3  *
4  * License:
5  *   This Source Code Form is subject to the terms of
6  *   the Mozilla Public License, v. 2.0. If a copy of
7  *   the MPL was not distributed with this file, You
8  *   can obtain one at http://mozilla.org/MPL/2.0/.
9  *
10  * Authors:
11  *   Vladimir Panteleev <ae@cy.md>
12  */
13 
14 module ae.sys.windows.process;
15 version (Windows):
16 
17 import std.exception;
18 import std.typecons;
19 
20 import ae.sys.windows.imports;
21 mixin(importWin32!q{w32api});
22 mixin(importWin32!q{winbase});
23 mixin(importWin32!q{windef});
24 mixin(importWin32!q{winuser});
25 
26 import ae.sys.windows.exception;
27 import ae.sys.windows.text;
28 
29 alias wenforce = ae.sys.windows.exception.wenforce;
30 
31 pragma(lib, "user32");
32 
33 static if (_WIN32_WINNT >= 0x500) {
34 
35 /// Wraps a created Windows process.
36 /// Similar to `std.process.Pid`
37 struct CreatedProcessImpl
38 {
39 	PROCESS_INFORMATION pi;
40 	alias pi this;
41 
42 	DWORD wait()
43 	{
44 		WaitForSingleObject(hProcess, INFINITE);
45 		DWORD dwExitCode;
46 		wenforce(GetExitCodeProcess(hProcess, &dwExitCode), "GetExitCodeProcess");
47 		return dwExitCode;
48 	}
49 
50 	~this()
51 	{
52 		CloseHandle(pi.hProcess);
53 		CloseHandle(pi.hThread);
54 	}
55 }
56 
57 alias RefCounted!CreatedProcessImpl CreatedProcess;
58 /// Create a Windows process.
59 CreatedProcess createProcess(string applicationName, string commandLine, STARTUPINFOW si = STARTUPINFOW.init)
60 {
61 	return createProcess(applicationName, commandLine, null, si);
62 }
63 
64 /// ditto
65 CreatedProcess createProcess(string applicationName, string commandLine, string currentDirectory, STARTUPINFOW si = STARTUPINFOW.init)
66 {
67 	CreatedProcess result;
68 	wenforce(CreateProcessW(toWStringz(applicationName), cast(LPWSTR)toWStringz(commandLine), null, null, false, 0, null, toWStringz(currentDirectory), &si, &result.pi), "CreateProcess");
69 	AllowSetForegroundWindow(result.dwProcessId);
70 	AttachThreadInput(GetCurrentThreadId(), result.dwThreadId, TRUE);
71 	AllowSetForegroundWindow(result.dwProcessId);
72 	return result;
73 }
74 
75 enum TOKEN_ADJUST_SESSIONID = 0x0100;
76 //enum SecurityImpersonation = 2;
77 //enum TokenPrimary = 1;
78 alias extern(Windows) BOOL function(
79   HANDLE hToken,
80   DWORD dwLogonFlags,
81   LPCWSTR lpApplicationName,
82   LPWSTR lpCommandLine,
83   DWORD dwCreationFlags,
84   LPVOID lpEnvironment,
85   LPCWSTR lpCurrentDirectory,
86   LPSTARTUPINFOW lpStartupInfo,
87   LPPROCESS_INFORMATION lpProcessInfo
88 ) CreateProcessWithTokenWFunc;
89 
90 /// Create a non-elevated process, if the current process is elevated.
91 CreatedProcess createDesktopUserProcess(string applicationName, string commandLine, STARTUPINFOW si = STARTUPINFOW.init)
92 {
93 	CreateProcessWithTokenWFunc CreateProcessWithTokenW = cast(CreateProcessWithTokenWFunc)GetProcAddress(GetModuleHandle("advapi32.dll"), "CreateProcessWithTokenW");
94 
95 	HANDLE hShellProcess = null, hShellProcessToken = null, hPrimaryToken = null;
96 	HWND hwnd = null;
97 	DWORD dwPID = 0;
98 
99 	// Enable SeIncreaseQuotaPrivilege in this process.  (This won't work if current process is not elevated.)
100 	HANDLE hProcessToken = null;
101 	wenforce(OpenProcessToken(GetCurrentProcess(), TOKEN_ADJUST_PRIVILEGES, &hProcessToken), "OpenProcessToken failed");
102 	scope(exit) CloseHandle(hProcessToken);
103 
104 	TOKEN_PRIVILEGES tkp;
105 	tkp.PrivilegeCount = 1;
106 	LookupPrivilegeValue(null, SE_INCREASE_QUOTA_NAME.ptr, &tkp.Privileges()[0].Luid);
107 	tkp.Privileges()[0].Attributes = SE_PRIVILEGE_ENABLED;
108 	wenforce(AdjustTokenPrivileges(hProcessToken, FALSE, &tkp, 0, null, null), "AdjustTokenPrivileges failed");
109 
110 	// Get an HWND representing the desktop shell.
111 	// CAVEATS:  This will fail if the shell is not running (crashed or terminated), or the default shell has been
112 	// replaced with a custom shell.  This also won't return what you probably want if Explorer has been terminated and
113 	// restarted elevated.
114 	hwnd = GetShellWindow();
115 	enforce(hwnd, "No desktop shell is present");
116 
117 	// Get the PID of the desktop shell process.
118 	GetWindowThreadProcessId(hwnd, &dwPID);
119 	enforce(dwPID, "Unable to get PID of desktop shell.");
120 
121 	// Open the desktop shell process in order to query it (get the token)
122 	hShellProcess = wenforce(OpenProcess(PROCESS_QUERY_INFORMATION, FALSE, dwPID), "Can't open desktop shell process");
123 	scope(exit) CloseHandle(hShellProcess);
124 
125 	// Get the process token of the desktop shell.
126 	wenforce(OpenProcessToken(hShellProcess, TOKEN_DUPLICATE, &hShellProcessToken), "Can't get process token of desktop shell");
127 	scope(exit) CloseHandle(hShellProcessToken);
128 
129 	// Duplicate the shell's process token to get a primary token.
130 	// Based on experimentation, this is the minimal set of rights required for CreateProcessWithTokenW (contrary to current documentation).
131 	const DWORD dwTokenRights = TOKEN_QUERY | TOKEN_ASSIGN_PRIMARY | TOKEN_DUPLICATE | TOKEN_ADJUST_DEFAULT | TOKEN_ADJUST_SESSIONID;
132 	wenforce(DuplicateTokenEx(hShellProcessToken, dwTokenRights, null, SECURITY_IMPERSONATION_LEVEL.SecurityImpersonation, TOKEN_TYPE.TokenPrimary, &hPrimaryToken), "Can't get primary token");
133 	scope(exit) CloseHandle(hPrimaryToken);
134 
135 	CreatedProcess result;
136 
137 	// Start the target process with the new token.
138 	wenforce(CreateProcessWithTokenW(
139 		hPrimaryToken,
140 		0,
141 		toWStringz(applicationName), cast(LPWSTR)toWStringz(commandLine),
142 		0,
143 		null,
144 		null,
145 		&si,
146 		&result.pi,
147 	), "CreateProcessWithTokenW failed");
148 
149 	AllowSetForegroundWindow(result.dwProcessId);
150 	AttachThreadInput(GetCurrentThreadId(), result.dwThreadId, TRUE);
151 	AllowSetForegroundWindow(result.dwProcessId);
152 
153 	return result;
154 }
155 
156 // --------------------------------------------------------------------------
157 
158 mixin(importWin32!q{tlhelp32});
159 
160 /// Wraps a `Toolhelp32Snapshot` handle.
161 struct ToolhelpSnapshotImpl
162 {
163 	HANDLE hSnapshot;
164 
165 	~this()
166 	{
167 		CloseHandle(hSnapshot);
168 	}
169 }
170 
171 alias RefCounted!ToolhelpSnapshotImpl ToolhelpSnapshot;
172 
173 /// `CreateToolhelp32Snapshot` wrapper.
174 ToolhelpSnapshot createToolhelpSnapshot(DWORD dwFlags, DWORD th32ProcessID=0)
175 {
176 	ToolhelpSnapshot result;
177 	auto hSnapshot = CreateToolhelp32Snapshot(dwFlags, th32ProcessID);
178 	wenforce(hSnapshot != INVALID_HANDLE_VALUE, "CreateToolhelp32Snapshot");
179 	result.hSnapshot = hSnapshot;
180 	return result;
181 }
182 
183 /// D ranges wrapping iteration over `Toolhelp32Snapshot` lists.
184 struct ToolhelpIterator(STRUCT, alias FirstFunc, alias NextFunc)
185 {
186 private:
187 	ToolhelpSnapshot snapshot;
188 	STRUCT s;
189 	BOOL bSuccess;
190 
191 	this(ToolhelpSnapshot snapshot)
192 	{
193 		this.snapshot = snapshot;
194 		s.dwSize = STRUCT.sizeof;
195 		bSuccess = FirstFunc(snapshot.hSnapshot, &s);
196 	}
197 
198 public:
199 	@property
200 	bool empty() const { return bSuccess == 0; } ///
201 
202 	@property
203 	ref STRUCT front() { return s; } ///
204 
205 	void popFront()
206 	{
207 		bSuccess = NextFunc(snapshot.hSnapshot, &s);
208 	} ///
209 }
210 
211 alias ProcessIterator = ToolhelpIterator!(PROCESSENTRY32, Process32First, Process32Next); /// ditto
212 @property ProcessIterator processes(ToolhelpSnapshot snapshot) { return ProcessIterator(snapshot); } /// ditto
213 
214 alias ThreadIterator = ToolhelpIterator!(THREADENTRY32, Thread32First, Thread32Next); /// ditto
215 @property ThreadIterator threads(ToolhelpSnapshot snapshot) { return ThreadIterator(snapshot); } /// ditto
216 
217 alias ModuleIterator = ToolhelpIterator!(MODULEENTRY32, Module32First, Module32Next); /// ditto
218 @property ModuleIterator modules(ToolhelpSnapshot snapshot) { return ModuleIterator(snapshot); } /// ditto
219 
220 alias HeapIterator = ToolhelpIterator!(HEAPLIST32, Heap32ListFirst, Heap32ListNext); /// ditto
221 @property HeapIterator heaps(ToolhelpSnapshot snapshot) { return HeapIterator(snapshot); } /// ditto
222 
223 // --------------------------------------------------------------------------
224 
225 /// Maintains a list of currently running processes,
226 /// and notifies the caller about new or exited processes.
227 struct ProcessWatcher
228 {
229 	PROCESSENTRY32[DWORD] oldProcesses; ///
230 
231 	void update(void delegate(ref PROCESSENTRY32) oldHandler, void delegate(ref PROCESSENTRY32) newHandler, bool handleExisting = false)
232 	{
233 		PROCESSENTRY32[DWORD] newProcesses;
234 		foreach (ref process; createToolhelpSnapshot(TH32CS_SNAPPROCESS).processes)
235 			newProcesses[process.th32ProcessID] = process;
236 
237 		if (oldProcesses || handleExisting) // Skip calling delegates on first run
238 		{
239 			if (oldHandler)
240 				foreach (pid, ref process; oldProcesses)
241 					if (pid !in newProcesses)
242 						oldHandler(process);
243 
244 			if (newHandler)
245 				foreach (pid, ref process; newProcesses)
246 					if (pid !in oldProcesses)
247 						newHandler(process);
248 		}
249 
250 		oldProcesses = newProcesses;
251 	} ///
252 }
253 
254 } // _WIN32_WINNT >= 0x500
255 
256 // ***************************************************************************
257 
258 /// Read/WriteProcessMemory helpers.
259 alias ubyte* RemoteAddress;
260 
261 void readProcessMemory(HANDLE h, RemoteAddress addr, void[] data)
262 {
263 	size_t c;
264 	wenforce(ReadProcessMemory(h, addr, data.ptr, data.length, &c), "ReadProcessMemory");
265 	enforce(c==data.length, "Not all data read");
266 } /// ditto
267 
268 void writeProcessMemory(HANDLE h, RemoteAddress addr, const(void)[] data)
269 {
270 	size_t c;
271 	wenforce(WriteProcessMemory(h, addr, data.ptr, data.length, &c), "WriteProcessMemory");
272 	enforce(c==data.length, "Not all data written");
273 } /// ditto
274 
275 void readProcessVar(T)(HANDLE h, RemoteAddress addr, T* v)
276 {
277 	h.readProcessMemory(addr, v[0..1]);
278 } /// ditto
279 
280 T readProcessVar(T)(HANDLE h, RemoteAddress addr)
281 {
282 	T v;
283 	h.readProcessVar(addr, &v);
284 	return v;
285 } /// ditto
286 
287 void writeProcessVar(T)(HANDLE h, RemoteAddress addr, auto ref T v)
288 {
289 	h.writeProcessMemory(addr, (&v)[0..1]);
290 } /// ditto
291 
292 /// Binding to a variable located in another process.
293 /// Automatically allocates and deallocates remote memory.
294 /// Use .read() and .write() to update local/remote data.
295 struct RemoteProcessVarImpl(T)
296 {
297 	T local; /// Cached local value.
298 	@property T* localPtr() { return &local; }
299 	RemoteAddress remotePtr; /// Address in remote process.
300 	HANDLE hProcess; /// Process handle.
301 
302 	this(HANDLE hProcess)
303 	{
304 		this.hProcess = hProcess;
305 		remotePtr = cast(RemoteAddress)wenforce(VirtualAllocEx(hProcess, null, T.sizeof, MEM_COMMIT, PAGE_READWRITE));
306 	} ///
307 
308 	void read()
309 	{
310 		readProcessMemory (hProcess, remotePtr, localPtr[0..1]);
311 	} ///
312 
313 	void write()
314 	{
315 		writeProcessMemory(hProcess, remotePtr, localPtr[0..1]);
316 	} ///
317 
318 	~this()
319 	{
320 		VirtualFreeEx(hProcess, remotePtr, 0, MEM_RELEASE);
321 	}
322 }
323 
324 template RemoteProcessVar(T)
325 {
326 	alias RefCounted!(RemoteProcessVarImpl!T) RemoteProcessVar;
327 } /// ditto