blob: fa33d4687f4c8b638b400e160c4f681fe8551ca0 [file] [log] [blame]
koder aka kdanilove06762a2015-03-22 23:32:09 +02001import re
koder aka kdanilov3a6633e2015-03-26 18:20:00 +02002import time
koder aka kdanilov416b87a2015-05-12 00:26:04 +03003import errno
koder aka kdanilov4af1c1d2015-05-18 15:48:58 +03004import random
koder aka kdanilov652cd802015-04-13 12:21:07 +03005import socket
koder aka kdanilov0c598a12015-04-21 03:01:40 +03006import shutil
koder aka kdanilove06762a2015-03-22 23:32:09 +02007import logging
8import os.path
koder aka kdanilov3a6633e2015-03-26 18:20:00 +02009import getpass
koder aka kdanilovf86d7af2015-05-06 04:01:54 +030010import StringIO
koder aka kdanilov652cd802015-04-13 12:21:07 +030011import threading
koder aka kdanilov0c598a12015-04-21 03:01:40 +030012import subprocess
koder aka kdanilov652cd802015-04-13 12:21:07 +030013
koder aka kdanilov3a6633e2015-03-26 18:20:00 +020014import paramiko
koder aka kdanilove06762a2015-03-22 23:32:09 +020015
koder aka kdanilove06762a2015-03-22 23:32:09 +020016
koder aka kdanilovcff7b2e2015-04-18 20:48:15 +030017logger = logging.getLogger("wally")
koder aka kdanilove06762a2015-03-22 23:32:09 +020018
19
koder aka kdanilov0c598a12015-04-21 03:01:40 +030020class Local(object):
21 "placeholder for local node"
22 @classmethod
23 def open_sftp(cls):
koder aka kdanilovafd98742015-04-24 01:27:22 +030024 return cls()
koder aka kdanilov0c598a12015-04-21 03:01:40 +030025
26 @classmethod
27 def mkdir(cls, remotepath, mode=None):
28 os.mkdir(remotepath)
29 if mode is not None:
30 os.chmod(remotepath, mode)
31
32 @classmethod
33 def put(cls, localfile, remfile):
koder aka kdanilov4d4771c2015-04-23 01:32:02 +030034 dirname = os.path.dirname(remfile)
35 if not os.path.exists(dirname):
36 os.makedirs(dirname)
koder aka kdanilov0c598a12015-04-21 03:01:40 +030037 shutil.copyfile(localfile, remfile)
38
39 @classmethod
koder aka kdanilova8253ce2015-06-13 03:10:21 +030040 def get(cls, remfile, localfile):
41 dirname = os.path.dirname(localfile)
42 if not os.path.exists(dirname):
43 os.makedirs(dirname)
44 shutil.copyfile(remfile, localfile)
45
46 @classmethod
koder aka kdanilov0c598a12015-04-21 03:01:40 +030047 def chmod(cls, path, mode):
48 os.chmod(path, mode)
49
50 @classmethod
51 def copytree(cls, src, dst):
52 shutil.copytree(src, dst)
53
54 @classmethod
55 def remove(cls, path):
56 os.unlink(path)
57
58 @classmethod
59 def close(cls):
60 pass
61
62 @classmethod
63 def open(cls, *args, **kwarhgs):
64 return open(*args, **kwarhgs)
65
koder aka kdanilove2de58c2015-04-24 22:59:36 +030066 @classmethod
67 def stat(cls, path):
68 return os.stat(path)
69
koder aka kdanilov783b4542015-04-23 18:57:04 +030070 def __enter__(self):
71 return self
72
73 def __exit__(self, x, y, z):
74 return False
75
koder aka kdanilov0c598a12015-04-21 03:01:40 +030076
koder aka kdanilovf86d7af2015-05-06 04:01:54 +030077NODE_KEYS = {}
78
79
koder aka kdanilov416b87a2015-05-12 00:26:04 +030080def exists(sftp, path):
81 """os.path.exists for paramiko's SCP object
82 """
83 try:
84 sftp.stat(path)
85 return True
86 except IOError as e:
87 if e.errno == errno.ENOENT:
88 return False
89 raise
90
91
koder aka kdanilovf86d7af2015-05-06 04:01:54 +030092def set_key_for_node(host_port, key):
93 sio = StringIO.StringIO(key)
94 NODE_KEYS[host_port] = paramiko.RSAKey.from_private_key(sio)
95 sio.close()
96
97
koder aka kdanilovbb5fe072015-05-21 02:50:23 +030098def ssh_connect(creds, conn_timeout=60, reuse_conn=None):
koder aka kdanilov0c598a12015-04-21 03:01:40 +030099 if creds == 'local':
koder aka kdanilov416b87a2015-05-12 00:26:04 +0300100 return Local()
koder aka kdanilov0c598a12015-04-21 03:01:40 +0300101
koder aka kdanilov46d4f392015-04-24 11:35:00 +0300102 tcp_timeout = 15
103 banner_timeout = 30
104
koder aka kdanilovbb5fe072015-05-21 02:50:23 +0300105 if reuse_conn is None:
106 ssh = paramiko.SSHClient()
107 ssh.load_host_keys('/dev/null')
108 ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
109 ssh.known_hosts = None
110 else:
111 ssh = reuse_conn
koder aka kdanilov168f6092015-04-19 02:33:38 +0300112
koder aka kdanilov6b1341a2015-04-21 22:44:21 +0300113 etime = time.time() + conn_timeout
114
115 while True:
koder aka kdanilov3a6633e2015-03-26 18:20:00 +0200116 try:
koder aka kdanilov46d4f392015-04-24 11:35:00 +0300117 tleft = etime - time.time()
118 c_tcp_timeout = min(tcp_timeout, tleft)
koder aka kdanilov8bc48022015-07-15 11:54:47 +0300119
120 if paramiko.__version_info__ >= (1, 15, 2):
121 banner_timeout = {'banner_timeout': min(banner_timeout, tleft)}
122 else:
123 banner_timeout = {}
koder aka kdanilov46d4f392015-04-24 11:35:00 +0300124
koder aka kdanilov3a6633e2015-03-26 18:20:00 +0200125 if creds.passwd is not None:
126 ssh.connect(creds.host,
koder aka kdanilov46d4f392015-04-24 11:35:00 +0300127 timeout=c_tcp_timeout,
koder aka kdanilova4a570f2015-04-23 22:11:40 +0300128 username=creds.user,
koder aka kdanilov3a6633e2015-03-26 18:20:00 +0200129 password=creds.passwd,
130 port=creds.port,
131 allow_agent=False,
koder aka kdanilov46d4f392015-04-24 11:35:00 +0300132 look_for_keys=False,
koder aka kdanilov8bc48022015-07-15 11:54:47 +0300133 **banner_timeout)
koder aka kdanilov46d4f392015-04-24 11:35:00 +0300134 elif creds.key_file is not None:
koder aka kdanilov3a6633e2015-03-26 18:20:00 +0200135 ssh.connect(creds.host,
koder aka kdanilova4a570f2015-04-23 22:11:40 +0300136 username=creds.user,
koder aka kdanilov46d4f392015-04-24 11:35:00 +0300137 timeout=c_tcp_timeout,
koder aka kdanilov3a6633e2015-03-26 18:20:00 +0200138 key_filename=creds.key_file,
139 look_for_keys=False,
koder aka kdanilov46d4f392015-04-24 11:35:00 +0300140 port=creds.port,
koder aka kdanilov8bc48022015-07-15 11:54:47 +0300141 **banner_timeout)
koder aka kdanilovf86d7af2015-05-06 04:01:54 +0300142 elif (creds.host, creds.port) in NODE_KEYS:
143 ssh.connect(creds.host,
144 username=creds.user,
145 timeout=c_tcp_timeout,
146 pkey=NODE_KEYS[(creds.host, creds.port)],
147 look_for_keys=False,
148 port=creds.port,
koder aka kdanilov8bc48022015-07-15 11:54:47 +0300149 **banner_timeout)
koder aka kdanilov46d4f392015-04-24 11:35:00 +0300150 else:
151 key_file = os.path.expanduser('~/.ssh/id_rsa')
152 ssh.connect(creds.host,
153 username=creds.user,
154 timeout=c_tcp_timeout,
155 key_filename=key_file,
156 look_for_keys=False,
157 port=creds.port,
koder aka kdanilov8bc48022015-07-15 11:54:47 +0300158 **banner_timeout)
koder aka kdanilov3a6633e2015-03-26 18:20:00 +0200159 return ssh
koder aka kdanilov3a6633e2015-03-26 18:20:00 +0200160 except paramiko.PasswordRequiredException:
161 raise
koder aka kdanilov46d4f392015-04-24 11:35:00 +0300162 except (socket.error, paramiko.SSHException):
koder aka kdanilov6b1341a2015-04-21 22:44:21 +0300163 if time.time() > etime:
koder aka kdanilov3a6633e2015-03-26 18:20:00 +0200164 raise
koder aka kdanilovcff7b2e2015-04-18 20:48:15 +0300165 time.sleep(1)
koder aka kdanilov3a6633e2015-03-26 18:20:00 +0200166
167
koder aka kdanilov4d4771c2015-04-23 01:32:02 +0300168def save_to_remote(sftp, path, content):
169 with sftp.open(path, "wb") as fd:
170 fd.write(content)
171
172
173def read_from_remote(sftp, path):
174 with sftp.open(path, "rb") as fd:
175 return fd.read()
176
177
koder aka kdanilove06762a2015-03-22 23:32:09 +0200178def normalize_dirpath(dirpath):
179 while dirpath.endswith("/"):
180 dirpath = dirpath[:-1]
181 return dirpath
182
183
koder aka kdanilov2c473092015-03-29 17:12:13 +0300184ALL_RWX_MODE = ((1 << 9) - 1)
185
186
187def ssh_mkdir(sftp, remotepath, mode=ALL_RWX_MODE, intermediate=False):
koder aka kdanilove06762a2015-03-22 23:32:09 +0200188 remotepath = normalize_dirpath(remotepath)
189 if intermediate:
190 try:
191 sftp.mkdir(remotepath, mode=mode)
koder aka kdanilov4d4771c2015-04-23 01:32:02 +0300192 except (IOError, OSError):
koder aka kdanilov168f6092015-04-19 02:33:38 +0300193 upper_dir = remotepath.rsplit("/", 1)[0]
194
195 if upper_dir == '' or upper_dir == '/':
196 raise
197
198 ssh_mkdir(sftp, upper_dir, mode=mode, intermediate=True)
koder aka kdanilove06762a2015-03-22 23:32:09 +0200199 return sftp.mkdir(remotepath, mode=mode)
200 else:
201 sftp.mkdir(remotepath, mode=mode)
202
203
204def ssh_copy_file(sftp, localfile, remfile, preserve_perm=True):
205 sftp.put(localfile, remfile)
206 if preserve_perm:
koder aka kdanilov2c473092015-03-29 17:12:13 +0300207 sftp.chmod(remfile, os.stat(localfile).st_mode & ALL_RWX_MODE)
koder aka kdanilove06762a2015-03-22 23:32:09 +0200208
209
210def put_dir_recursively(sftp, localpath, remotepath, preserve_perm=True):
211 "upload local directory to remote recursively"
212
213 # hack for localhost connection
214 if hasattr(sftp, "copytree"):
215 sftp.copytree(localpath, remotepath)
216 return
217
218 assert remotepath.startswith("/"), "%s must be absolute path" % remotepath
219
220 # normalize
221 localpath = normalize_dirpath(localpath)
222 remotepath = normalize_dirpath(remotepath)
223
224 try:
225 sftp.chdir(remotepath)
226 localsuffix = localpath.rsplit("/", 1)[1]
227 remotesuffix = remotepath.rsplit("/", 1)[1]
228 if localsuffix != remotesuffix:
229 remotepath = os.path.join(remotepath, localsuffix)
230 except IOError:
231 pass
232
233 for root, dirs, fls in os.walk(localpath):
234 prefix = os.path.commonprefix([localpath, root])
235 suffix = root.split(prefix, 1)[1]
236 if suffix.startswith("/"):
237 suffix = suffix[1:]
238
239 remroot = os.path.join(remotepath, suffix)
240
241 try:
242 sftp.chdir(remroot)
243 except IOError:
244 if preserve_perm:
koder aka kdanilov2c473092015-03-29 17:12:13 +0300245 mode = os.stat(root).st_mode & ALL_RWX_MODE
koder aka kdanilove06762a2015-03-22 23:32:09 +0200246 else:
koder aka kdanilov2c473092015-03-29 17:12:13 +0300247 mode = ALL_RWX_MODE
koder aka kdanilove06762a2015-03-22 23:32:09 +0200248 ssh_mkdir(sftp, remroot, mode=mode, intermediate=True)
249 sftp.chdir(remroot)
250
251 for f in fls:
252 remfile = os.path.join(remroot, f)
253 localfile = os.path.join(root, f)
254 ssh_copy_file(sftp, localfile, remfile, preserve_perm)
255
256
koder aka kdanilov4500a5f2015-04-17 16:55:17 +0300257def delete_file(conn, path):
258 sftp = conn.open_sftp()
259 sftp.remove(path)
260 sftp.close()
261
262
koder aka kdanilove06762a2015-03-22 23:32:09 +0200263def copy_paths(conn, paths):
264 sftp = conn.open_sftp()
265 try:
266 for src, dst in paths.items():
267 try:
268 if os.path.isfile(src):
269 ssh_copy_file(sftp, src, dst)
270 elif os.path.isdir(src):
271 put_dir_recursively(sftp, src, dst)
272 else:
273 templ = "Can't copy {0!r} - " + \
274 "it neither a file not a directory"
koder aka kdanilov168f6092015-04-19 02:33:38 +0300275 raise OSError(templ.format(src))
koder aka kdanilove06762a2015-03-22 23:32:09 +0200276 except Exception as exc:
277 tmpl = "Scp {0!r} => {1!r} failed - {2!r}"
koder aka kdanilov168f6092015-04-19 02:33:38 +0300278 raise OSError(tmpl.format(src, dst, exc))
koder aka kdanilove06762a2015-03-22 23:32:09 +0200279 finally:
280 sftp.close()
281
282
283class ConnCreds(object):
koder aka kdanilov2c473092015-03-29 17:12:13 +0300284 conn_uri_attrs = ("user", "passwd", "host", "port", "path")
285
koder aka kdanilove06762a2015-03-22 23:32:09 +0200286 def __init__(self):
koder aka kdanilov2c473092015-03-29 17:12:13 +0300287 for name in self.conn_uri_attrs:
koder aka kdanilove06762a2015-03-22 23:32:09 +0200288 setattr(self, name, None)
289
koder aka kdanilovcff7b2e2015-04-18 20:48:15 +0300290 def __str__(self):
291 return str(self.__dict__)
292
koder aka kdanilove06762a2015-03-22 23:32:09 +0200293
294uri_reg_exprs = []
295
296
297class URIsNamespace(object):
298 class ReParts(object):
299 user_rr = "[^:]*?"
koder aka kdanilov7e0f7cf2015-05-01 17:24:35 +0300300 host_rr = "[^:@]*?"
koder aka kdanilove06762a2015-03-22 23:32:09 +0200301 port_rr = "\\d+"
302 key_file_rr = "[^:@]*"
303 passwd_rr = ".*?"
304
305 re_dct = ReParts.__dict__
306
307 for attr_name, val in re_dct.items():
308 if attr_name.endswith('_rr'):
309 new_rr = "(?P<{0}>{1})".format(attr_name[:-3], val)
310 setattr(ReParts, attr_name, new_rr)
311
312 re_dct = ReParts.__dict__
313
314 templs = [
315 "^{host_rr}$",
koder aka kdanilov7e0f7cf2015-05-01 17:24:35 +0300316 "^{host_rr}:{port_rr}$",
koder aka kdanilov416b87a2015-05-12 00:26:04 +0300317 "^{host_rr}::{key_file_rr}$",
318 "^{host_rr}:{port_rr}:{key_file_rr}$",
koder aka kdanilov7e0f7cf2015-05-01 17:24:35 +0300319 "^{user_rr}@{host_rr}$",
320 "^{user_rr}@{host_rr}:{port_rr}$",
koder aka kdanilove06762a2015-03-22 23:32:09 +0200321 "^{user_rr}@{host_rr}::{key_file_rr}$",
322 "^{user_rr}@{host_rr}:{port_rr}:{key_file_rr}$",
koder aka kdanilov7e0f7cf2015-05-01 17:24:35 +0300323 "^{user_rr}:{passwd_rr}@{host_rr}$",
324 "^{user_rr}:{passwd_rr}@{host_rr}:{port_rr}$",
koder aka kdanilove06762a2015-03-22 23:32:09 +0200325 ]
326
327 for templ in templs:
328 uri_reg_exprs.append(templ.format(**re_dct))
329
330
331def parse_ssh_uri(uri):
koder aka kdanilov7e0f7cf2015-05-01 17:24:35 +0300332 # user:passwd@ip_host:port
333 # user:passwd@ip_host
koder aka kdanilove06762a2015-03-22 23:32:09 +0200334 # user@ip_host:port
335 # user@ip_host
336 # ip_host:port
337 # ip_host
338 # user@ip_host:port:path_to_key_file
339 # user@ip_host::path_to_key_file
340 # ip_host:port:path_to_key_file
341 # ip_host::path_to_key_file
342
koder aka kdanilov4d4771c2015-04-23 01:32:02 +0300343 if uri.startswith("ssh://"):
344 uri = uri[len("ssh://"):]
345
koder aka kdanilove06762a2015-03-22 23:32:09 +0200346 res = ConnCreds()
347 res.port = "22"
348 res.key_file = None
349 res.passwd = None
koder aka kdanilova4a570f2015-04-23 22:11:40 +0300350 res.user = getpass.getuser()
koder aka kdanilove06762a2015-03-22 23:32:09 +0200351
352 for rr in uri_reg_exprs:
353 rrm = re.match(rr, uri)
354 if rrm is not None:
355 res.__dict__.update(rrm.groupdict())
356 return res
koder aka kdanilov652cd802015-04-13 12:21:07 +0300357
koder aka kdanilove06762a2015-03-22 23:32:09 +0200358 raise ValueError("Can't parse {0!r} as ssh uri value".format(uri))
359
360
koder aka kdanilovbb5fe072015-05-21 02:50:23 +0300361def reconnect(conn, uri, **params):
362 if uri == 'local':
363 return conn
364
365 creds = parse_ssh_uri(uri)
366 creds.port = int(creds.port)
367 return ssh_connect(creds, reuse_conn=conn, **params)
368
369
koder aka kdanilov168f6092015-04-19 02:33:38 +0300370def connect(uri, **params):
koder aka kdanilov0c598a12015-04-21 03:01:40 +0300371 if uri == 'local':
koder aka kdanilov0fdaaee2015-06-30 11:10:48 +0300372 res = Local()
373 else:
374 creds = parse_ssh_uri(uri)
375 creds.port = int(creds.port)
376 res = ssh_connect(creds, **params)
377 return res
koder aka kdanilove06762a2015-03-22 23:32:09 +0200378
379
koder aka kdanilov652cd802015-04-13 12:21:07 +0300380all_sessions_lock = threading.Lock()
koder aka kdanilov416b87a2015-05-12 00:26:04 +0300381all_sessions = {}
382
383
koder aka kdanilov4af1c1d2015-05-18 15:48:58 +0300384class BGSSHTask(object):
385 def __init__(self, node, use_sudo):
386 self.node = node
387 self.pid = None
388 self.use_sudo = use_sudo
koder aka kdanilov416b87a2015-05-12 00:26:04 +0300389
koder aka kdanilov4af1c1d2015-05-18 15:48:58 +0300390 def start(self, orig_cmd, **params):
391 uniq_name = 'test'
392 cmd = "screen -S {0} -d -m {1}".format(uniq_name, orig_cmd)
393 run_over_ssh(self.node.connection, cmd,
394 timeout=10, node=self.node.get_conn_id(),
395 **params)
396 processes = run_over_ssh(self.node.connection, "ps aux", nolog=True)
koder aka kdanilov416b87a2015-05-12 00:26:04 +0300397
koder aka kdanilov4af1c1d2015-05-18 15:48:58 +0300398 for proc in processes.split("\n"):
399 if orig_cmd in proc and "SCREEN" not in proc:
400 self.pid = proc.split()[1]
401 break
402 else:
403 self.pid = -1
404
405 def check_running(self):
406 assert self.pid is not None
407 try:
408 run_over_ssh(self.node.connection,
409 "ls /proc/{0}".format(self.pid),
410 timeout=10, nolog=True)
411 return True
412 except OSError:
413 return False
koder aka kdanilov4af1c1d2015-05-18 15:48:58 +0300414
415 def kill(self, soft=True, use_sudo=True):
416 assert self.pid is not None
417 try:
418 if soft:
419 cmd = "kill {0}"
420 else:
421 cmd = "kill -9 {0}"
422
423 if self.use_sudo:
424 cmd = "sudo " + cmd
425
426 run_over_ssh(self.node.connection,
427 cmd.format(self.pid), nolog=True)
428 return True
429 except OSError:
430 return False
431
432 def wait(self, soft_timeout, timeout):
433 end_of_wait_time = timeout + time.time()
434 soft_end_of_wait_time = soft_timeout + time.time()
koder aka kdanilov4af1c1d2015-05-18 15:48:58 +0300435
koder aka kdanilova8253ce2015-06-13 03:10:21 +0300436 # time_till_check = random.randint(5, 10)
437 time_till_check = 2
438
439 # time_till_first_check = random.randint(2, 6)
440 time_till_first_check = 2
koder aka kdanilovfd2cfa52015-05-20 03:17:42 +0300441 time.sleep(time_till_first_check)
442 if not self.check_running():
443 return True
444
koder aka kdanilov4af1c1d2015-05-18 15:48:58 +0300445 while self.check_running() and time.time() < soft_end_of_wait_time:
446 time.sleep(soft_end_of_wait_time - time.time())
447
448 while end_of_wait_time > time.time():
449 time.sleep(time_till_check)
450 if not self.check_running():
451 break
452 else:
453 self.kill()
koder aka kdanilova8253ce2015-06-13 03:10:21 +0300454 time.sleep(1)
koder aka kdanilov4af1c1d2015-05-18 15:48:58 +0300455 if self.check_running():
456 self.kill(soft=False)
457 return False
458 return True
koder aka kdanilove06762a2015-03-22 23:32:09 +0200459
koder aka kdanilove06762a2015-03-22 23:32:09 +0200460
koder aka kdanilovcff7b2e2015-04-18 20:48:15 +0300461def run_over_ssh(conn, cmd, stdin_data=None, timeout=60,
462 nolog=False, node=None):
koder aka kdanilov652cd802015-04-13 12:21:07 +0300463 "should be replaces by normal implementation, with select"
koder aka kdanilov0c598a12015-04-21 03:01:40 +0300464
koder aka kdanilov416b87a2015-05-12 00:26:04 +0300465 if isinstance(conn, Local):
koder aka kdanilov0c598a12015-04-21 03:01:40 +0300466 if not nolog:
467 logger.debug("SSH:local Exec {0!r}".format(cmd))
468 proc = subprocess.Popen(cmd, shell=True,
469 stdin=subprocess.PIPE,
470 stdout=subprocess.PIPE,
471 stderr=subprocess.STDOUT)
472
473 stdoutdata, _ = proc.communicate(input=stdin_data)
koder aka kdanilov0c598a12015-04-21 03:01:40 +0300474 if proc.returncode != 0:
475 templ = "SSH:{0} Cmd {1!r} failed with code {2}. Output: {3}"
476 raise OSError(templ.format(node, cmd, proc.returncode, stdoutdata))
477
478 return stdoutdata
479
koder aka kdanilov652cd802015-04-13 12:21:07 +0300480 transport = conn.get_transport()
481 session = transport.open_session()
koder aka kdanilove06762a2015-03-22 23:32:09 +0200482
koder aka kdanilov4500a5f2015-04-17 16:55:17 +0300483 if node is None:
484 node = ""
485
koder aka kdanilov652cd802015-04-13 12:21:07 +0300486 with all_sessions_lock:
koder aka kdanilov416b87a2015-05-12 00:26:04 +0300487 all_sessions[id(session)] = session
koder aka kdanilove06762a2015-03-22 23:32:09 +0200488
koder aka kdanilov652cd802015-04-13 12:21:07 +0300489 try:
490 session.set_combine_stderr(True)
koder aka kdanilove06762a2015-03-22 23:32:09 +0200491
koder aka kdanilov652cd802015-04-13 12:21:07 +0300492 stime = time.time()
koder aka kdanilove06762a2015-03-22 23:32:09 +0200493
koder aka kdanilov652cd802015-04-13 12:21:07 +0300494 if not nolog:
koder aka kdanilov4500a5f2015-04-17 16:55:17 +0300495 logger.debug("SSH:{0} Exec {1!r}".format(node, cmd))
koder aka kdanilove06762a2015-03-22 23:32:09 +0200496
koder aka kdanilov652cd802015-04-13 12:21:07 +0300497 session.exec_command(cmd)
koder aka kdanilove06762a2015-03-22 23:32:09 +0200498
koder aka kdanilov652cd802015-04-13 12:21:07 +0300499 if stdin_data is not None:
500 session.sendall(stdin_data)
koder aka kdanilove06762a2015-03-22 23:32:09 +0200501
koder aka kdanilov652cd802015-04-13 12:21:07 +0300502 session.settimeout(1)
503 session.shutdown_write()
504 output = ""
505
506 while True:
507 try:
508 ndata = session.recv(1024)
509 output += ndata
510 if "" == ndata:
511 break
512 except socket.timeout:
513 pass
514
515 if time.time() - stime > timeout:
516 raise OSError(output + "\nExecution timeout")
517
518 code = session.recv_exit_status()
519 finally:
koder aka kdanilov4af1c1d2015-05-18 15:48:58 +0300520 found = False
koder aka kdanilov416b87a2015-05-12 00:26:04 +0300521 with all_sessions_lock:
koder aka kdanilov4af1c1d2015-05-18 15:48:58 +0300522 if id(session) in all_sessions:
523 found = True
524 del all_sessions[id(session)]
525
526 if found:
527 session.close()
koder aka kdanilov652cd802015-04-13 12:21:07 +0300528
529 if code != 0:
koder aka kdanilov4500a5f2015-04-17 16:55:17 +0300530 templ = "SSH:{0} Cmd {1!r} failed with code {2}. Output: {3}"
531 raise OSError(templ.format(node, cmd, code, output))
koder aka kdanilov652cd802015-04-13 12:21:07 +0300532
533 return output
534
535
536def close_all_sessions():
537 with all_sessions_lock:
koder aka kdanilov416b87a2015-05-12 00:26:04 +0300538 for session in all_sessions.values():
koder aka kdanilov652cd802015-04-13 12:21:07 +0300539 try:
540 session.sendall('\x03')
541 session.close()
542 except:
543 pass
koder aka kdanilov416b87a2015-05-12 00:26:04 +0300544 all_sessions.clear()