blob: bdfe6b56a9113ffbf110b888ee24b8f3d107dd18 [file] [log] [blame]
Alex9a4ad212020-10-01 18:04:25 -05001import queue
2import subprocess
3import traceback
4import threading
5
6from time import sleep
7from .exception import TimeoutException, CheckerException
8from .other import shell, piped_shell
9from .log import logger, logger_cli
10
11
12# We do not use paramiko here to preserve system level ssh config
13def ssh_shell_p(
14 command,
15 host,
16 username=None,
17 keypath=None,
18 port=None,
19 silent=False,
20 piped=False,
21 use_sudo=False
22):
23 _ssh_cmd = []
24 _ssh_cmd.append("ssh")
25 if silent:
26 _ssh_cmd.append("-q")
27 # Build SSH cmd
28 if keypath:
29 _ssh_cmd.append("-i " + keypath)
Alexc4f59622021-08-27 13:42:00 -050030 _ssh_cmd.append("-o " + "IdentitiesOnly=yes")
Alex9a4ad212020-10-01 18:04:25 -050031 if port:
32 _ssh_cmd.append("-p " + str(port))
33 if username:
34 _ssh_cmd.append(username+'@'+host)
35 else:
36 _ssh_cmd.append(host)
37
38 if use_sudo:
39 _ssh_cmd.append("sudo")
40
41 _ssh_cmd.append(command)
42
43 _ssh_cmd = " ".join(_ssh_cmd)
44 if not piped:
45 return shell(_ssh_cmd)
46 else:
47 return piped_shell(_ssh_cmd)
48
49
50def scp_p(
51 source,
52 target,
53 port=None,
54 keypath=None,
55 silent=False,
56 piped=False
57):
58 _scp_cmd = []
59 _scp_cmd.append("scp")
60 if port:
61 _scp_cmd.append("-P " + str(port))
62 if silent:
63 _scp_cmd.append("-q")
64 # Build SSH cmd
65 if keypath:
66 _scp_cmd.append("-i " + keypath)
67 _scp_cmd.append(source)
68 _scp_cmd.append(target)
69 _scp_cmd = " ".join(_scp_cmd)
70 if not piped:
71 return shell(_scp_cmd)
72 else:
73 return piped_shell(_scp_cmd)
74
75
76def output_reader(_stdout, outq):
77 for line in iter(_stdout.readline, b''):
78 outq.put(line.decode('utf-8'))
79
80
81# Base class for all SSH related actions
82class SshBase(object):
83 def __init__(
84 self,
85 tgt_host,
86 user=None,
87 keypath=None,
88 port=None,
89 timeout=15,
90 silent=False,
91 piped=False
92 ):
93 self._cmd = ["ssh"]
94 self.timeout = timeout
95 self.port = port if port else 22
96 self.host = tgt_host
97 self.username = user
98 self.keypath = keypath
99 self.silent = silent
100 self.piped = piped
101 self.output = []
102
103 self._options = ["-tt"]
104 if self.keypath:
105 self._options += ["-i", self.keypath]
106 if self.port:
107 self._options += ["-p", str(self.port)]
108 self._extra_options = [
109 "-o", "UserKnownHostsFile=/dev/null",
Alexc4f59622021-08-27 13:42:00 -0500110 "-o", "StrictHostKeyChecking=no",
111 "-o", "IdentitiesOnly=yes"
Alex9a4ad212020-10-01 18:04:25 -0500112 ]
113
114 self._host_uri = ""
115 if self.username:
116 self._host_uri = self.username + "@" + self.host
117 else:
118 self._host_uri = self.host
119
120 def _connect(self, banner="Welcome"):
121 if not isinstance(banner, str):
122 raise CheckerException(
123 "Invalid SSH banner type: '{}'".format(type(banner))
124 )
Alexc4f59622021-08-27 13:42:00 -0500125 logger.debug("... connecting")
Alex9a4ad212020-10-01 18:04:25 -0500126 while True:
127 try:
128 line = self.outq.get(block=False)
129 self.output.append(line)
130 if line.startswith(banner):
131 break
132 except queue.Empty:
133 logger.debug("... {} sec".format(self.timeout))
134 sleep(1)
135 self.timeout -= 1
136 if not self.timeout:
137 logger.debug(
138 "...timed out after {} sec".format(str(self.timeout))
139 )
140 return False
Alexc4f59622021-08-27 13:42:00 -0500141 logger.debug("... connected")
Alex9a4ad212020-10-01 18:04:25 -0500142 return True
143
144 def _wait_for_string(self, string):
Alexc4f59622021-08-27 13:42:00 -0500145 logger.debug("... waiting for '{}'".format(string))
Alex9a4ad212020-10-01 18:04:25 -0500146 while True:
147 try:
148 line = self.outq.get(block=False)
149 line = line.decode() if isinstance(line, bytes) else line
150 self.output.append(line)
151 if not line.startswith(string):
152 continue
153 else:
154 break
155 except queue.Empty:
156 logger.debug("... {} sec".format(self.timeout))
157 sleep(1)
158 self.timeout -= 1
159 if not self.timeout:
160 logger.debug(
Alexc4f59622021-08-27 13:42:00 -0500161 "... timed out after {} sec".format(str(self.timeout))
Alex9a4ad212020-10-01 18:04:25 -0500162 )
163 return False
Alexc4f59622021-08-27 13:42:00 -0500164 logger.debug("... found")
Alex9a4ad212020-10-01 18:04:25 -0500165 return True
166
167 def _init_connection(self, cmd):
168 self._proc = subprocess.Popen(
169 cmd,
170 stdin=subprocess.PIPE,
171 stdout=subprocess.PIPE,
172 stderr=subprocess.PIPE,
173 universal_newlines=False,
174 bufsize=0
175 )
176 # Create thread safe output getter
177 self.outq = queue.Queue()
178 self._t = threading.Thread(
179 target=output_reader,
180 args=(self._proc.stdout, self.outq)
181 )
182 self._t.start()
183
184 # Track if there is an yes/no
185 if not self._connect():
186 raise TimeoutException(
187 "SSH connection to '{}'".format(self.host)
188 )
189
190 self.input = self._proc.stdin
191 self.get_output()
192 logger.debug(
193 "Connected. Banners:\n{}".format(
194 "".join(self.flush_output())
195 )
196 )
197
198 def _end_connection(self):
199 # Kill the ssh process if it is alive
200 if self._proc.poll() is None:
201 self._proc.kill()
202 self.get_output()
203
204 return
205
206 def do(self, cmd, timeout=30, sudo=False, strip_cmd=True):
207 cmd = cmd if isinstance(cmd, bytes) else bytes(cmd.encode('utf-8'))
Alexc4f59622021-08-27 13:42:00 -0500208 logger.debug("... ssh: '{}'".format(cmd))
Alex9a4ad212020-10-01 18:04:25 -0500209 if sudo:
210 _cmd = b"sudo " + cmd
211 else:
212 _cmd = cmd
213 # run command
214 self.input.write(_cmd + b'\n')
215 # wait for completion
216 self.wait_ready(_cmd, timeout=timeout)
217 self.get_output()
218 _output = self.flush_output().replace('\r', '')
219 if strip_cmd:
220 return "\n".join(_output.splitlines()[1:])
221 else:
222 return _output
223
224 def get_output(self):
225 while True:
226 try:
227 line = self.outq.get(block=False)
228 line = str(line) if isinstance(line, bytes) else line
229 self.output.append(line)
230 except queue.Empty:
231 return self.output
232 return None
233
234 def flush_output(self, as_string=True):
235 _out = self.output
236 self.output = []
237 if as_string:
238 return "".join(_out)
239 else:
240 return _out
241
242 def wait_ready(self, cmd, timeout=60):
243 def _strip_cmd_carrets(_str, carret='\r', skip_chars=1):
244 _cnt = _str.count(carret)
245 while _cnt > 0:
246 _idx = _str.index(carret)
247 _str = _str[:_idx] + _str[_idx+1+skip_chars:]
248 _cnt -= 1
249 return _str
250 while True:
251 try:
252 _line = self.outq.get(block=False)
253 line = _line.decode() if isinstance(_line, bytes) else _line
254 # line = line.replace('\r', '')
255 self.output.append(line)
256 # check if this is the command itself and skip
257 if '$' in line:
258 _cmd = line.split('$', 1)[1].strip()
259 _cmd = _strip_cmd_carrets(_cmd)
260 if _cmd == cmd.decode():
261 continue
262 break
263 except queue.Empty:
264 logger.debug("... {} sec".format(timeout))
265 sleep(1)
266 timeout -= 1
267 if not timeout:
268 logger.debug("...timed out")
269 return False
270 return True
271
272 def wait_for_string(self, string, timeout=60):
273 if not self._wait_for_string(string):
274 raise TimeoutException(
275 "Time out waiting for string '{}'".format(string)
276 )
277 else:
278 return True
279
280
281class SshShell(SshBase):
282 def __enter__(self):
283 self._cmd = ["ssh"]
284 self._cmd += self._options
285 self._cmd += self._extra_options
286 self._cmd += [self._host_uri]
287
288 logger.debug("...shell to: '{}'".format(" ".join(self._cmd)))
289 self._init_connection(self._cmd)
290 return self
291
292 def __exit__(self, _type, _value, _traceback):
293 self._end_connection()
294 if _value:
295 logger.warn(
296 "Error running SSH:\r\n{}".format(
297 "".join(traceback.format_exception(
298 _type,
299 _value,
300 _traceback
301 ))
302 )
303 )
304
305 return True
306
307 def connect(self):
308 return self.__enter__()
309
310 def kill(self):
311 self._end_connection()
312
313 def get_host_path(self, path):
314 _uri = self.host + ":" + path
315 if self.username:
316 _uri = self.username + "@" + _uri
317 return _uri
318
319 def scp(self, _src, _dst):
320 self._scp_options = []
321 if self.keypath:
322 self._scp_options += ["-i", self.keypath]
323 if self.port:
324 self._scp_options += ["-P", str(self.port)]
325
326 _cmd = ["scp"]
327 _cmd += self._scp_options
328 _cmd += self._extra_options
329 _cmd += [_src]
330 _cmd += [_dst]
331
332 logger.debug("...scp: '{}'".format(" ".join(_cmd)))
333 _proc = subprocess.Popen(
334 _cmd,
335 stdout=subprocess.PIPE,
336 stderr=subprocess.PIPE
337 )
338 _r = _proc.communicate()
339 _e = _r[1].decode() if _r[1] else ""
340 return _proc.returncode, _r[0].decode(), _e
341
342
343class PortForward(SshBase):
344 def __init__(
345 self,
346 host,
347 fwd_host,
348 user=None,
349 keypath=None,
350 port=None,
351 loc_port=10022,
352 fwd_port=22,
353 timeout=15
354 ):
355 super(PortForward, self).__init__(
356 host,
357 user=user,
358 keypath=keypath,
359 port=port,
360 timeout=timeout,
361 silent=True,
362 piped=False
363 )
364 self.f_host = fwd_host
365 self.l_port = loc_port
366 self.f_port = fwd_port
367
368 self._forward_options = [
369 "-L",
370 ":".join([str(self.l_port), self.f_host, str(self.f_port)])
371 ]
372
373 def __enter__(self):
374 self._cmd = ["ssh"]
375 self._cmd += self._forward_options
376 self._cmd += self._options
377 self._cmd += self._extra_options
378 self._cmd += [self._host_uri]
379
380 logger.debug(
Alexc4f59622021-08-27 13:42:00 -0500381 "... port forwarding: '{}'".format(" ".join(self._cmd))
Alex9a4ad212020-10-01 18:04:25 -0500382 )
383 self._init_connection(self._cmd)
384 return self
385
386 def __exit__(self, _type, _value, _traceback):
387 self._end_connection()
388 if _value:
389 logger_cli.warn(
390 "Error running SSH:\r\n{}".format(
391 "".join(traceback.format_exception(
392 _type,
393 _value,
394 _traceback
395 ))
396 )
397 )
398
399 return True
400
401 def connect(self):
402 return self.__enter__()
403
404 def kill(self):
405 self._end_connection()