From 1f5edecbf6dadb72936923cf2532be97c27b9f3e Mon Sep 17 00:00:00 2001
From: Akash Mozumdar <akashmozumdar@gmail.com>
Date: Mon, 18 Feb 2019 22:39:04 -0500
Subject: [PATCH] attach to command line args

---
 GUI/host/util.cpp  |  9 +++++++++
 GUI/host/util.h    |  1 +
 GUI/mainwindow.cpp | 20 +++++++++++++++-----
 3 files changed, 25 insertions(+), 5 deletions(-)

diff --git a/GUI/host/util.cpp b/GUI/host/util.cpp
index 69dec83..7df766c 100644
--- a/GUI/host/util.cpp
+++ b/GUI/host/util.cpp
@@ -265,6 +265,15 @@ namespace Util
 		return {};
 	}
 
+	std::vector<DWORD> GetAllProcessIds()
+	{
+		std::vector<DWORD> processIds(10000);
+		DWORD spaceUsed = 0;
+		EnumProcesses(processIds.data(), 10000 * sizeof(DWORD), &spaceUsed);
+		processIds.resize(spaceUsed / sizeof(DWORD));
+		return processIds;
+	}
+
 	std::optional<std::wstring> GetClipboardText()
 	{
 		if (!IsClipboardFormatAvailable(CF_UNICODETEXT)) return {};
diff --git a/GUI/host/util.h b/GUI/host/util.h
index b4ef1e2..a9e309b 100644
--- a/GUI/host/util.h
+++ b/GUI/host/util.h
@@ -7,6 +7,7 @@ namespace Util
 {
 	std::optional<std::wstring> GetModuleFilename(DWORD processId, HMODULE module = NULL);
 	std::optional<std::wstring> GetModuleFilename(HMODULE module = NULL);
+	std::vector<DWORD> GetAllProcessIds();
 	std::optional<std::wstring> GetClipboardText();
 	std::optional<std::wstring> StringToWideString(const std::string& text, UINT encoding = CP_UTF8);
 	// return true if repetition found (see https://github.com/Artikash/Textractor/issues/40)
diff --git a/GUI/mainwindow.cpp b/GUI/mainwindow.cpp
index 2889dbb..fa9e697 100644
--- a/GUI/mainwindow.cpp
+++ b/GUI/mainwindow.cpp
@@ -5,7 +5,7 @@
 #include "extenwindow.h"
 #include "misc.h"
 #include "host/util.h"
-#include <Psapi.h>
+#include <shellapi.h>
 #include <winhttp.h>
 #include <QFormLayout>
 #include <QPushButton>
@@ -58,6 +58,18 @@ MainWindow::MainWindow(QWidget *parent) :
 	current = &Host::GetThread(Host::console);
 	Host::AddConsoleOutput(ABOUT);
 
+	std::vector<DWORD> processIds = Util::GetAllProcessIds();
+	std::vector<std::wstring> processNames;
+	for (auto processId : processIds) processNames.emplace_back(Util::GetModuleFilename(processId).value_or(L""));
+	int argc;
+	std::unique_ptr<LPWSTR[], Functor<LocalFree>> argv(CommandLineToArgvW(GetCommandLineW(), &argc));
+	for (int i = 0; i < argc; ++i)
+		if (std::wstring arg = argv[i]; arg[0] == L'/' || arg[0] == L'-')
+			if (arg[1] == L'P')
+				if (DWORD processId = _wtoi(arg.substr(2).c_str())) Host::InjectProcess(processId);
+				else for (int i = 0; i < processIds.size(); ++i)
+					if (processNames[i].find(L"\\" + arg.substr(2)) != std::wstring::npos) Host::InjectProcess(processIds[i]);
+
 	std::thread([]
 	{
 		using InternetHandle = AutoHandle<Functor<WinHttpCloseHandle>>;
@@ -193,10 +205,8 @@ std::array<InfoForExtension, 10> MainWindow::GetMiscInfo(TextThread& thread)
 void MainWindow::AttachProcess()
 {
 	QMultiHash<QString, DWORD> allProcesses;
-	DWORD allProcessIds[5000] = {}, spaceUsed = 0;
-	EnumProcesses(allProcessIds, sizeof(allProcessIds), &spaceUsed);
-	for (int i = 0; i < spaceUsed / sizeof(DWORD); ++i)
-		if (auto processName = Util::GetModuleFilename(allProcessIds[i])) allProcesses.insert(QFileInfo(S(processName.value())).fileName(), allProcessIds[i]);
+	for (auto processId : Util::GetAllProcessIds())
+		if (auto processName = Util::GetModuleFilename(processId)) allProcesses.insert(QFileInfo(S(processName.value())).fileName(), processId);
 
 	QStringList processList(allProcesses.uniqueKeys());
 	processList.sort(Qt::CaseInsensitive);