diff --git a/src/agent/Agent.cc b/src/agent/Agent.cc index efc3e7f..a77b644 100644 --- a/src/agent/Agent.cc +++ b/src/agent/Agent.cc @@ -332,6 +332,9 @@ void Agent::handlePacket(ReadBuffer &packet) // at once, we can ignore the early ones. handleSetSizePacket(packet); break; + case AgentMsg::GetConsoleProcessList: + handleGetConsoleProcessListPacket(packet); + break; default: trace("Unrecognized message, id:%d", type); } @@ -426,6 +429,29 @@ void Agent::handleSetSizePacket(ReadBuffer &packet) writePacket(reply); } +void Agent::handleGetConsoleProcessListPacket(ReadBuffer &packet) +{ + packet.assertEof(); + + auto processList = std::vector(64); + auto processCount = GetConsoleProcessList(&processList[0], processList.size()); + if (processList.size() < processCount) { + processList.resize(processCount); + processCount = GetConsoleProcessList(&processList[0], processList.size()); + } + + if (processCount == 0) { + trace("GetConsoleProcessList failed"); + } + + auto reply = newPacket(); + reply.putInt32(processCount); + for (DWORD i = 0; i < processCount; i++) { + reply.putInt32(processList[i]); + } + writePacket(reply); +} + void Agent::pollConinPipe() { const std::string newData = m_coninPipe->readAllToString(); diff --git a/src/agent/Agent.h b/src/agent/Agent.h index 5775e08..1dde48f 100644 --- a/src/agent/Agent.h +++ b/src/agent/Agent.h @@ -59,6 +59,7 @@ private: void writePacket(WriteBuffer &packet); void handleStartProcessPacket(ReadBuffer &packet); void handleSetSizePacket(ReadBuffer &packet); + void handleGetConsoleProcessListPacket(ReadBuffer &packet); void pollConinPipe(); protected: diff --git a/src/include/winpty.h b/src/include/winpty.h index 2a196a5..fdfe4bc 100644 --- a/src/include/winpty.h +++ b/src/include/winpty.h @@ -218,6 +218,11 @@ WINPTY_API BOOL winpty_set_size(winpty_t *wp, int cols, int rows, winpty_error_ptr_t *err /*OPTIONAL*/); +/* Gets a list of processes attached to the console. */ +WINPTY_API int +winpty_get_console_process_list(winpty_t *wp, int *processList, const int processCount, + winpty_error_ptr_t *err /*OPTIONAL*/); + /* Frees the winpty_t object and the OS resources contained in it. This * call breaks the connection with the agent, which should then close its * console, terminating the processes attached to it. diff --git a/src/libwinpty/winpty.cc b/src/libwinpty/winpty.cc index e89e348..3d97749 100644 --- a/src/libwinpty/winpty.cc +++ b/src/libwinpty/winpty.cc @@ -935,6 +935,33 @@ winpty_set_size(winpty_t *wp, int cols, int rows, } API_CATCH(FALSE) } +WINPTY_API int +winpty_get_console_process_list(winpty_t *wp, int *processList, const int processCount, + winpty_error_ptr_t *err /*OPTIONAL*/) { + API_TRY { + ASSERT(wp != nullptr); + ASSERT(processList != nullptr); + LockGuard lock(wp->mutex); + RpcOperation rpc(*wp); + auto packet = newPacket(); + packet.putInt32(AgentMsg::GetConsoleProcessList); + writePacket(*wp, packet); + auto reply = readPacket(*wp); + + auto actualProcessCount = reply.getInt32(); + + if (actualProcessCount <= processCount) { + for (auto i = 0; i < actualProcessCount; i++) { + processList[i] = reply.getInt32(); + } + } + + reply.assertEof(); + rpc.success(); + return actualProcessCount; + } API_CATCH(0) +} + WINPTY_API void winpty_free(winpty_t *wp) { // At least in principle, CloseHandle can fail, so this deletion can // fail. It won't throw an exception, but maybe there's an error that diff --git a/src/shared/AgentMsg.h b/src/shared/AgentMsg.h index 31859e0..ab60c6b 100644 --- a/src/shared/AgentMsg.h +++ b/src/shared/AgentMsg.h @@ -26,6 +26,7 @@ struct AgentMsg enum Type { StartProcess, SetSize, + GetConsoleProcessList, }; };