AdbTools/CommandExecution.cpp at master · WangONC/AdbTools · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
#include "pch.h"
#include "CommandExecution.h"
#include <memory>
#include <thread>
#include <tchar.h>

CommandExecution::CommandResult CommandExecution::ExecCommand(const CString& command, bool asAdmin)
{
    return ExecCommand(std::vector<CString>{command}, asAdmin);
}

CommandExecution::CommandResult CommandExecution::ExecCommand(const std::vector<CString>& commands, bool asAdmin)
{
    if (commands.empty()) {
        return CommandResult();
    }

    CString commandLine = _T("powershell.exe -NoProfile -ExecutionPolicy Bypass -Command \"");
    for (const auto& command : commands) {
        commandLine += EscapePowerShellCommand(command) + _T("; ");
    }
    commandLine += _T("\"");

    return ExecuteCommands(commandLine, asAdmin);
}

std::future<CommandExecution::CommandResult> CommandExecution::ExecCommandAsync(const CString& command, bool asAdmin)
{
    return ExecCommandAsync(std::vector<CString>{command}, asAdmin);
}

std::future<CommandExecution::CommandResult> CommandExecution::ExecCommandAsync(const std::vector<CString>& commands, bool asAdmin)
{
    if (commands.empty()) {
        std::promise<CommandResult> resultPromise;
        resultPromise.set_value(CommandResult());
        return resultPromise.get_future();
    }

    CString commandLine = _T("powershell.exe -NoProfile -ExecutionPolicy Bypass -Command \"");
    for (const auto& command : commands) {
        commandLine += EscapePowerShellCommand(command) + _T("; ");
    }
    commandLine += _T("\"");

    std::promise<CommandResult> resultPromise;
    std::future<CommandResult> resultFuture = resultPromise.get_future();

    std::thread([commandLine, asAdmin, resultPromise = std::move(resultPromise)]() mutable {
        AsyncExecutor(commandLine, asAdmin, std::move(resultPromise));
        }).detach();

    return resultFuture;
}

CommandExecution::CommandResult CommandExecution::ExecuteCommands(const CString& commandLine, bool asAdmin)
{
    CommandResult commandResult;

    SECURITY_ATTRIBUTES securityAttributes = { sizeof(SECURITY_ATTRIBUTES), nullptr, TRUE };
    HANDLE hReadPipe = nullptr, hWritePipe = nullptr;
    if (!CreatePipe(&hReadPipe, &hWritePipe, &securityAttributes, 0)) {
        commandResult.errorMsg = _T("Failed to create pipe");
        return commandResult;
    }

    STARTUPINFO startupInfo = { sizeof(STARTUPINFO) };
    startupInfo.dwFlags = STARTF_USESHOWWINDOW | STARTF_USESTDHANDLES;
    startupInfo.wShowWindow = SW_HIDE;
    startupInfo.hStdOutput = hWritePipe;
    startupInfo.hStdError = hWritePipe;

    PROCESS_INFORMATION processInfo = { 0 };

    TCHAR* commandLineCopy = _tcsdup(commandLine);

    DWORD creationFlags = asAdmin ? CREATE_NEW_CONSOLE : CREATE_NO_WINDOW;
    BOOL success = CreateProcess(
        nullptr,
        commandLineCopy,
        nullptr,
        nullptr,
        TRUE,
        creationFlags,
        nullptr,
        nullptr,
        &startupInfo,
        &processInfo
    );

    if (!success) {
        commandResult.errorMsg = _T("Failed to create process");
        CloseHandle(hReadPipe);
        CloseHandle(hWritePipe);
        return commandResult;
    }

    CloseHandle(hWritePipe);

    const int bufferSize = 4096;
    std::unique_ptr<CHAR[]> buffer(new CHAR[bufferSize]);
    DWORD bytesRead;
    CString output;

    while (ReadFile(hReadPipe, buffer.get(), bufferSize - 1, &bytesRead, nullptr) && bytesRead > 0) {
        buffer[bytesRead] = '\0';
        output += CString(buffer.get());
    }

    WaitForSingleObject(processInfo.hProcess, INFINITE);
    GetExitCodeProcess(processInfo.hProcess, &commandResult.result);

    CloseHandle(processInfo.hProcess);
    CloseHandle(processInfo.hThread);
    CloseHandle(hReadPipe);

    if (commandResult.result == 0) commandResult.successMsg = output;
    else commandResult.errorMsg = output;

    ATLTRACE(_T("Result: %d | Success: %s | Error: %s\n"),
        commandResult.result,
        commandResult.successMsg,
        commandResult.errorMsg);

    return commandResult;
}

void CommandExecution::AsyncExecutor(const CString& commandLine, bool asAdmin, std::promise<CommandResult>&& resultPromise)
{
    CommandResult result = ExecuteCommands(commandLine, asAdmin);
    resultPromise.set_value(std::move(result));
}

CString CommandExecution::EscapePowerShellCommand(const CString& command)
{
    CString escaped = command;
    escaped.Replace(_T("\""), _T("\\\"")); // 转义双引号
    //escaped.Replace(_T("$"), _T("`$")); // 转义美元符号 这里的操作如果是直接在终端中执行确实是需要的，但是通过“运行”的方式，也就是直接创建进程的方式，是不需要的
    //return _T("\"") + escaped + _T("\"");
    return escaped;
}