Package starcluster :: Module ssh
[hide private]
[frames] | no frames]

Source Code for Module starcluster.ssh

  1  """ 
  2  ssh.py 
  3  Friendly Python SSH2 interface. 
  4  From http://commandline.org.uk/code/ 
  5  License: LGPL 
  6  modified by justin riley (justin.t.riley@gmail.com) 
  7  """ 
  8   
  9  import os 
 10  import re 
 11  import sys 
 12  import stat 
 13  import string 
 14  import socket 
 15  import paramiko 
 16  import posixpath 
 17   
 18  # windows does not have termios... 
 19  try: 
 20      import termios 
 21      import tty 
 22      HAS_TERMIOS = True 
 23  except ImportError: 
 24      HAS_TERMIOS = False 
 25   
 26  from starcluster import exception 
 27  from starcluster.logger import log 
28 29 30 -class SSHClient(object):
31 """ 32 Establishes an SSH connection to a remote host using either password or 33 private key authentication. Once established, this object allows executing 34 commands, copying files to/from the remote host, various file querying 35 similar to os.path.*, and much more. 36 """ 37
38 - def __init__(self, 39 host, 40 username=None, 41 password=None, 42 private_key=None, 43 private_key_pass=None, 44 port=22, 45 timeout=30):
46 self._host = host 47 self._port = 22 48 self._pkey = None 49 self._username = username or os.environ['LOGNAME'] 50 self._password = password 51 self._timeout = timeout 52 self._sftp = None 53 self._transport = None 54 if private_key: 55 self._pkey = self.load_private_key(private_key, private_key_pass) 56 elif not password: 57 raise exception.SSHNoCredentialsError()
58
59 - def load_private_key(self, private_key, private_key_pass=None):
60 # Use Private Key. 61 log.debug('loading private key %s' % private_key) 62 if private_key.endswith('rsa') or private_key.count('rsa'): 63 pkey = self._load_rsa_key(private_key, private_key_pass) 64 elif private_key.endswith('dsa') or private_key.count('dsa'): 65 pkey = self._load_dsa_key(private_key, private_key_pass) 66 else: 67 log.debug("specified key does not end in either rsa or dsa" + \ 68 ", trying both") 69 pkey = self._load_rsa_key(private_key, private_key_pass) 70 if pkey is None: 71 pkey = self._load_dsa_key(private_key, private_key_pass) 72 return pkey
73
74 - def connect(self, host=None, username=None, password=None, 75 private_key=None, private_key_pass=None, port=22, timeout=30):
76 host = host or self._host 77 username = username or self._username 78 pkey = self._pkey 79 if private_key: 80 pkey = self.load_private_key(private_key, private_key_pass) 81 log.debug("connecting to host %s on port %d as user %s" % (host, port, 82 username)) 83 try: 84 sock = self._get_socket(host, port) 85 transport = paramiko.Transport(sock) 86 transport.banner_timeout = timeout 87 except socket.error: 88 raise exception.SSHConnectionError(host, port) 89 # Authenticate the transport. 90 try: 91 transport.connect(username=username, pkey=pkey, password=password) 92 except paramiko.AuthenticationException: 93 raise exception.SSHAuthException(username, host) 94 except paramiko.SSHException, e: 95 msg = e.args[0] 96 raise exception.SSHError(msg) 97 except socket.error: 98 raise exception.SSHConnectionError(host, port) 99 except EOFError: 100 raise exception.SSHConnectionError(host, port) 101 except Exception, e: 102 raise exception.SSHError(str(e)) 103 self.close() 104 self._transport = transport 105 return self
106 107 @property
108 - def transport(self):
109 """ 110 This property attempts to return an active SSH transport 111 """ 112 if not self._transport or not self._transport.is_active(): 113 self.connect(self._host, self._username, self._password, 114 port=self._port, timeout=self._timeout) 115 return self._transport
116
117 - def get_server_public_key(self):
118 return self.transport.get_remote_server_key()
119
120 - def is_active(self):
121 if self._transport: 122 return self._transport.is_active() 123 return False
124
125 - def _get_socket(self, hostname, port):
126 for (family, socktype, proto, canonname, sockaddr) in \ 127 socket.getaddrinfo(hostname, port, socket.AF_UNSPEC, 128 socket.SOCK_STREAM): 129 if socktype == socket.SOCK_STREAM: 130 af = family 131 break 132 else: 133 raise exception.SSHError( 134 'No suitable address family for %s' % hostname) 135 sock = socket.socket(af, socket.SOCK_STREAM) 136 sock.settimeout(self._timeout) 137 sock.connect((hostname, port)) 138 return sock
139
140 - def _load_rsa_key(self, private_key, private_key_pass=None):
141 private_key_file = os.path.expanduser(private_key) 142 try: 143 rsa_key = paramiko.RSAKey.from_private_key_file(private_key_file, 144 private_key_pass) 145 log.debug("Using private key %s (rsa)" % private_key) 146 return rsa_key 147 except paramiko.SSHException: 148 log.error('invalid rsa key or passphrase specified')
149
150 - def _load_dsa_key(self, private_key, private_key_pass=None):
151 private_key_file = os.path.expanduser(private_key) 152 try: 153 dsa_key = paramiko.DSSKey.from_private_key_file(private_key_file, 154 private_key_pass) 155 log.info("Using private key %s (dsa)" % private_key) 156 return dsa_key 157 except paramiko.SSHException: 158 log.error('invalid dsa key or passphrase specified')
159 160 @property
161 - def sftp(self):
162 """Establish the SFTP connection.""" 163 if not self._sftp or self._sftp.sock.closed: 164 log.debug("creating sftp connection") 165 self._sftp = paramiko.SFTPClient.from_transport(self.transport) 166 return self._sftp
167
168 - def generate_rsa_key(self):
169 return paramiko.RSAKey.generate(2048)
170
171 - def get_public_key(self, key):
172 return ' '.join([key.get_name(), key.get_base64()])
173
174 - def load_remote_rsa_key(self, remote_filename):
175 """ 176 Returns paramiko.RSAKey object for an RSA key located on the remote 177 machine 178 """ 179 rfile = self.remote_file(remote_filename, 'r') 180 key = paramiko.RSAKey(file_obj=rfile) 181 rfile.close() 182 return key
183
184 - def makedirs(self, path, mode=0755):
185 """ 186 Same as os.makedirs - makes a new directory and automatically creates 187 all parent directories if they do not exist. 188 189 mode specifies unix permissions to apply to the new dir 190 """ 191 head, tail = posixpath.split(path) 192 if not tail: 193 head, tail = posixpath.split(head) 194 if head and tail and not self.path_exists(head): 195 try: 196 self.makedirs(head, mode) 197 except OSError, e: 198 # be happy if someone already created the path 199 if e.errno != os.errno.EEXIST: 200 raise 201 # xxx/newdir/. exists if xxx/newdir exists 202 if tail == posixpath.curdir: 203 return 204 self.mkdir(path, mode)
205
206 - def mkdir(self, path, mode=0755, ignore_failure=False):
207 """ 208 Make a new directory on the remote machine 209 210 If parent is True, create all parent directories that do not exist 211 212 mode specifies unix permissions to apply to the new dir 213 """ 214 try: 215 return self.sftp.mkdir(path, mode) 216 except IOError: 217 if not ignore_failure: 218 raise
219
220 - def get_remote_file_lines(self, remote_file, regex=None, matching=True):
221 """ 222 Returns list of lines in a remote_file 223 224 If regex is passed only lines that contain a pattern that matches 225 regex will be returned 226 227 If matching is set to False then only lines *not* containing a pattern 228 that matches regex will be returned 229 """ 230 f = self.remote_file(remote_file, 'r') 231 flines = f.readlines() 232 f.close() 233 if regex is None: 234 return flines 235 r = re.compile(regex) 236 lines = [] 237 for line in flines: 238 match = r.search(line) 239 if matching and match: 240 lines.append(line) 241 elif not matching and not match: 242 lines.append(line) 243 return lines
244
245 - def remove_lines_from_file(self, remote_file, regex):
246 """ 247 Removes lines matching regex from remote_file 248 """ 249 if regex in [None, '']: 250 log.debug('no regex supplied...returning') 251 return 252 lines = self.get_remote_file_lines(remote_file, regex, matching=False) 253 log.debug("new %s after removing regex (%s) matches:\n%s" % \ 254 (remote_file, regex, ''.join(lines))) 255 f = self.remote_file(remote_file) 256 f.writelines(lines) 257 f.close()
258 261
262 - def remote_file(self, file, mode='w'):
263 """ 264 Returns a remote file descriptor 265 """ 266 rfile = self.sftp.open(file, mode) 267 rfile.name = file 268 return rfile
269
270 - def path_exists(self, path):
271 """ 272 Test whether a remote path exists. 273 Returns False for broken symbolic links 274 """ 275 try: 276 self.stat(path) 277 return True 278 except IOError: 279 return False
280
281 - def chown(self, uid, gid, remote_file):
282 """ 283 Apply permissions (mode) to remote_file 284 """ 285 f = self.remote_file(remote_file, 'r') 286 f.chown(uid, gid, remote_file) 287 f.close()
288
289 - def chmod(self, mode, remote_file):
290 """ 291 Apply permissions (mode) to remote_file 292 """ 293 f = self.remote_file(remote_file, 'r') 294 f.chmod(mode) 295 f.close()
296
297 - def ls(self, path):
298 """ 299 Return a list containing the names of the entries in the remote path. 300 """ 301 return [os.path.join(path, f) for f in self.sftp.listdir(path)]
302
303 - def isdir(self, path):
304 """ 305 Return true if the remote path refers to an existing directory. 306 """ 307 try: 308 s = self.stat(path) 309 except IOError: 310 return False 311 return stat.S_ISDIR(s.st_mode)
312
313 - def isfile(self, path):
314 """ 315 Return true if the remote path refers to an existing file. 316 """ 317 try: 318 s = self.stat(path) 319 except IOError: 320 return False 321 return stat.S_ISREG(s.st_mode)
322
323 - def stat(self, path):
324 """ 325 Perform a stat system call on the given remote path. 326 """ 327 return self.sftp.stat(path)
328
329 - def get(self, remotepath, localpath=None):
330 """ 331 Copies a file between the remote host and the local host. 332 """ 333 if not localpath: 334 localpath = os.path.split(remotepath)[1] 335 self.sftp_connect() 336 self.sftp.get(remotepath, localpath)
337
338 - def put(self, localpath, remotepath=None):
339 """ 340 Copies a file between the local host and the remote host. 341 """ 342 if not remotepath: 343 remotepath = os.path.split(localpath)[1] 344 self.sftp.put(localpath, remotepath)
345
346 - def execute_async(self, command):
347 """ 348 Executes a remote command without blocking 349 350 NOTE: this method will not block, however, if your process does not 351 complete or background itself before the python process executing this 352 code exits, it will not persist on the remote machine 353 """ 354 355 channel = self.transport.open_session() 356 channel.exec_command(command)
357
358 - def execute(self, command, silent=True, only_printable=False, 359 ignore_exit_status=False, log_output=True):
360 """ 361 Execute a remote command and return stdout/stderr 362 363 NOTE: this function blocks until the process finishes 364 365 kwargs: 366 silent - do not print output 367 only_printable - filter the command's output to allow only printable 368 characters 369 returns List of output lines 370 """ 371 channel = self.transport.open_session() 372 channel.exec_command(command) 373 #stdin = channel.makefile('wb', -1) 374 stdout = channel.makefile('rb', -1) 375 stderr = channel.makefile_stderr('rb', -1) 376 output = [] 377 line = None 378 if silent: 379 output = stdout.readlines() + stderr.readlines() 380 else: 381 while line != '': 382 line = stdout.readline() 383 if only_printable: 384 line = ''.join(c for c in line if c in string.printable) 385 if line != '': 386 output.append(line) 387 print line, 388 for line in stderr.readlines(): 389 output.append(line) 390 print line 391 if only_printable: 392 output = map(lambda line: ''.join(c for c in line if c in 393 string.printable), output) 394 output = map(lambda line: line.strip(), output) 395 exit_status = channel.recv_exit_status() 396 if exit_status != 0: 397 if not ignore_exit_status: 398 log.error("command '%s' failed with status %d" % (command, 399 exit_status)) 400 else: 401 log.debug("command %s failed with status %d" % (command, 402 exit_status)) 403 if log_output: 404 for line in output: 405 log.debug(line.strip()) 406 return output
407
408 - def has_required(self, progs):
409 """ 410 Same as check_required but returns False if not all commands exist 411 """ 412 try: 413 return self.check_required(progs) 414 except exception.RemoteCommandNotFound: 415 return False
416
417 - def check_required(self, progs):
418 """ 419 Checks that all commands in the progs list exist on the remote system. 420 Returns True if all commands exist and raises exception.CommandNotFound 421 if not. 422 """ 423 for prog in progs: 424 if not self.which(prog): 425 raise exception.RemoteCommandNotFound(prog) 426 return True
427
428 - def which(self, prog):
429 return self.execute('which %s' % prog, ignore_exit_status=True)
430
431 - def get_path(self):
432 """Returns the PATH environment variable on the remote machine""" 433 return self.get_env()['PATH']
434
435 - def get_env(self):
436 """Returns the remote machine's environment as a dictionary""" 437 env = {} 438 for line in self.execute('env'): 439 key, val = line.split('=', 1) 440 env[key] = val 441 return env
442
443 - def close(self):
444 """Closes the connection and cleans up.""" 445 if self._sftp: 446 self._sftp.close() 447 if self._transport: 448 self._transport.close()
449
450 - def interactive_shell(self, user='root'):
451 if user and self.transport.get_username() != user: 452 self.connect(username=user) 453 try: 454 chan = self.transport.open_session() 455 chan.get_pty() 456 chan.invoke_shell() 457 log.info('Starting interactive shell...') 458 if HAS_TERMIOS: 459 self._posix_shell(chan) 460 else: 461 self._windows_shell(chan) 462 chan.close() 463 except Exception, e: 464 import traceback 465 print '*** Caught exception: %s: %s' % (e.__class__, e) 466 traceback.print_exc()
467
468 - def _posix_shell(self, chan):
469 import select 470 471 oldtty = termios.tcgetattr(sys.stdin) 472 try: 473 tty.setraw(sys.stdin.fileno()) 474 tty.setcbreak(sys.stdin.fileno()) 475 chan.settimeout(0.0) 476 477 # needs to be sent to give vim correct size FIX 478 chan.send('eval $(resize)\n') 479 480 while True: 481 r, w, e = select.select([chan, sys.stdin], [], []) 482 if chan in r: 483 try: 484 x = chan.recv(1024) 485 if len(x) == 0: 486 print '\r\n*** EOF\r\n', 487 break 488 sys.stdout.write(x) 489 sys.stdout.flush() 490 except socket.timeout: 491 pass 492 if sys.stdin in r: 493 # fixes up arrow problem 494 x = os.read(sys.stdin.fileno(), 1) 495 if len(x) == 0: 496 break 497 chan.send(x) 498 finally: 499 termios.tcsetattr(sys.stdin, termios.TCSADRAIN, oldtty)
500 501 # thanks to Mike Looijmans for this code
502 - def _windows_shell(self, chan):
503 import threading 504 505 sys.stdout.write("Line-buffered terminal emulation. " + \ 506 "Press F6 or ^Z to send EOF.\r\n\r\n") 507 508 def writeall(sock): 509 while True: 510 data = sock.recv(256) 511 if not data: 512 sys.stdout.write('\r\n*** EOF ***\r\n\r\n') 513 sys.stdout.flush() 514 break 515 sys.stdout.write(data) 516 sys.stdout.flush()
517 518 writer = threading.Thread(target=writeall, args=(chan,)) 519 writer.start() 520 521 # needs to be sent to give vim correct size FIX 522 chan.send('eval $(resize)\n') 523 524 try: 525 while True: 526 d = sys.stdin.read(1) 527 if not d: 528 break 529 chan.send(d) 530 except EOFError: 531 # user hit ^Z or F6 532 pass
533
534 - def __del__(self):
535 """Attempt to clean up if not explicitly closed.""" 536 log.debug('__del__ called') 537 self.close()
538 539 540 # for backwards compatibility 541 Connection = SSHClient
542 543 544 -def main():
545 """Little test when called directly.""" 546 # Set these to your own details. 547 myssh = SSHClient('somehost.domain.com') 548 print myssh.execute('hostname') 549 #myssh.put('ssh.py') 550 myssh.close()
551 552 if __name__ == "__main__": 553 main() 554