1 package org.metricshub.wmi.remotecommand;
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23 import java.util.ArrayList;
24 import java.util.Collections;
25 import java.util.HashMap;
26 import java.util.List;
27 import java.util.Map;
28 import java.util.Objects;
29 import java.util.concurrent.TimeoutException;
30 import java.util.stream.Collectors;
31 import org.metricshub.wmi.TimeoutHelper;
32 import org.metricshub.wmi.Utils;
33 import org.metricshub.wmi.WmiHelper;
34 import org.metricshub.wmi.exceptions.ProcessNotFoundException;
35 import org.metricshub.wmi.exceptions.WmiComException;
36 import org.metricshub.wmi.exceptions.WqlQuerySyntaxException;
37 import org.metricshub.wmi.wbem.WmiWbemServices;
38
39
40
41
42
43 public class RemoteProcess {
44
45 private RemoteProcess() {}
46
47 private static final String TERMINATE = "Terminate";
48 private static final String CREATE = "Create";
49
50 private static final String CIMV2_NAMESPACE = "ROOT\\CIMV2";
51
52 private static final String WIN32_PROCESS = "Win32_Process";
53
54
55
56
57 private static final Map<Integer, String> METHOD_RETURNVALUE_MAP;
58
59 static {
60 final Map<Integer, String> map = new HashMap<>();
61 map.put(2, "Access denied");
62 map.put(3, "Insufficient privilege");
63 map.put(8, "Unknown failure");
64 map.put(9, "Path not found");
65 map.put(21, "Invalid parameter");
66 METHOD_RETURNVALUE_MAP = Collections.unmodifiableMap(map);
67 }
68
69
70
71
72
73
74
75
76
77
78
79
80
81 public static int executeCommand(
82 final String command,
83 final String hostname,
84 final String username,
85 final char[] password,
86 final String workingDirectory,
87 final long timeout
88 ) throws WmiComException, TimeoutException {
89 Utils.checkNonNull(command, "command");
90 Utils.checkArgumentNotZeroOrNegative(timeout, "timeout");
91
92 final long start = Utils.getCurrentTimeMillis();
93
94 final String networkResource = WmiHelper.createNetworkResource(hostname, CIMV2_NAMESPACE);
95 try (WmiWbemServices wmiWbemServices = WmiWbemServices.getInstance(networkResource, username, password)) {
96
97 final Map<String, Object> createInputs = new HashMap<>();
98 createInputs.put("CommandLine", command);
99 if (!Utils.isBlank(workingDirectory)) {
100 createInputs.put("CurrentDirectory", workingDirectory.trim());
101 }
102 final Map<String, Object> createResult = wmiWbemServices.executeMethod(
103 WIN32_PROCESS,
104 WIN32_PROCESS,
105 CREATE,
106 createInputs
107 );
108
109
110 final Integer processId = (Integer) createResult.get("ProcessId");
111 if (processId == null || processId.intValue() < 1) {
112 throw new WmiComException("Could not spawn the process: No ProcessId was returned by Win32_Process::Create");
113 }
114
115
116 try {
117 while (
118 existProcess(
119 wmiWbemServices,
120 processId,
121 TimeoutHelper.getRemainingTime(timeout, start, "No time left to check if the process exists")
122 )
123 ) {
124 TimeoutHelper.stagedSleep(timeout, start, String.format("Command %s execution has timed out", command));
125 }
126 } catch (final TimeoutException e) {
127
128
129 killProcessWithChildren(wmiWbemServices, processId, 10000);
130 throw e;
131 }
132
133 return (Integer) createResult.get("ReturnValue");
134 }
135 }
136
137
138
139
140
141
142
143
144
145
146
147
148 static boolean existProcess(final WmiWbemServices wbemServices, final int pid, final long timeout)
149 throws WmiComException, TimeoutException {
150 try {
151 return !wbemServices
152 .executeWql(String.format("SELECT Handle FROM Win32_Process WHERE Handle = '%d'", pid), timeout)
153 .isEmpty();
154 } catch (final WqlQuerySyntaxException e) {
155 throw new WmiComException(e);
156 }
157 }
158
159
160
161
162
163
164
165
166
167
168 private static void killProcessWithChildren(final WmiWbemServices wmiWbemServices, final int pid, final long timeout)
169 throws WmiComException, TimeoutException {
170
171 try {
172 final long start = Utils.getCurrentTimeMillis();
173
174 final List<Integer> pidToKillList = new ArrayList<>();
175 pidToKillList.add(pid);
176 pidToKillList.addAll(
177 wmiWbemServices
178 .executeWql(String.format("SELECT Handle FROM Win32_Process WHERE ParentProcessId = '%d'", pid), timeout)
179 .stream()
180 .map(row -> (String) row.get("Handle"))
181 .filter(Objects::nonNull)
182 .map(Integer::parseInt)
183 .collect(Collectors.toList())
184 );
185
186
187 for (final int pidToKill : pidToKillList) {
188 if (TimeoutHelper.getRemainingTime(timeout, start, "No time left to kill the process") < 0) {
189 throw new TimeoutException("Timeout while killing remaining processes");
190 }
191 try {
192 killProcess(wmiWbemServices, pidToKill);
193 } catch (final ProcessNotFoundException e) {
194
195 }
196 }
197 } catch (final WqlQuerySyntaxException e) {
198 throw new WmiComException(e);
199 }
200 }
201
202
203
204
205
206
207
208
209 private static void killProcess(final WmiWbemServices wmiWbemServices, final int pid)
210 throws WmiComException, ProcessNotFoundException {
211
212
213 final Map<String, Object> inputs = Collections.singletonMap("Reason", 1);
214 final Map<String, Object> terminateResult;
215 try {
216 terminateResult =
217 wmiWbemServices.executeMethod(
218 String.format("Win32_Process.Handle='%d'", pid),
219 WIN32_PROCESS,
220 TERMINATE,
221 inputs
222 );
223 } catch (final WmiComException e) {
224
225 if (e.getMessage().contains("WBEM_E_NOT_FOUND")) {
226 throw new ProcessNotFoundException(pid);
227 }
228 throw e;
229 }
230
231
232 final Integer returnCode = (Integer) terminateResult.get("ReturnValue");
233 if (returnCode == null || returnCode.intValue() != 0) {
234 throw new WmiComException("Could not terminate the process (%d): %s", pid, getReturnErrorMessage(returnCode));
235 }
236 }
237
238 private static String getReturnErrorMessage(final int returnCode) {
239 return METHOD_RETURNVALUE_MAP.getOrDefault(returnCode, String.format("Unknown return code (%d)", returnCode));
240 }
241 }