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 stat 
 11  import string 
 12  import tempfile 
 13  import paramiko 
 14   
 15  import socket 
 16  import sys 
 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 -class Connection(object):
30 """Establishes an ssh connection to a remote host using either password or private 31 key authentication. Once established, this object allows executing and 32 copying files to/from the remote host""" 33
34 - def __init__(self, 35 host, 36 username = None, 37 password = None, 38 private_key = None, 39 private_key_pass = None, 40 port = 22, 41 timeout=30, 42 ):
43 self._timeout = timeout 44 self._sftp_live = False 45 self._sftp = None 46 if not username: 47 username = os.environ['LOGNAME'] 48 49 # Log to a temporary file. 50 templog = tempfile.mkstemp('.txt', 'ssh-')[1] 51 paramiko.util.log_to_file(templog) 52 53 # Begin the SSH transport. 54 self._transport_live = False 55 try: 56 sock = self._get_socket(host, port) 57 self._transport = paramiko.Transport(sock) 58 except socket.error: 59 raise exception.SSHConnectionError(host, port) 60 self._transport_live = True 61 # Authenticate the transport. 62 if password: 63 # Using Password. 64 try: 65 self._transport.connect(username = username, password = password) 66 except paramiko.AuthenticationException: 67 raise exception.SSHAuthException(username,host) 68 elif private_key: 69 # Use Private Key. 70 pkey = None 71 log.debug('private key specified') 72 if private_key.endswith('rsa') or private_key.count('rsa'): 73 pkey = self._load_rsa_key(private_key, private_key_pass) 74 elif private_key.endswith('dsa') or private_key.count('dsa'): 75 pkey = self._load_dsa_key(private_key, private_key_pass) 76 else: 77 log.warn("specified key does not end in either rsa or dsa, trying both") 78 pkey = self._load_rsa_key(private_key, private_key_pass) 79 if pkey is None: 80 pkey = self._load_dsa_key(private_key, private_key_pass) 81 try: 82 self._transport.connect(username = username, pkey = pkey) 83 except paramiko.AuthenticationException: 84 raise exception.SSHAuthException(username, host) 85 else: 86 raise exception.SSHNoCredentialsError()
87
88 - def _get_socket(self, hostname, port):
89 for (family, socktype, proto, canonname, sockaddr) in \ 90 socket.getaddrinfo(hostname, port, socket.AF_UNSPEC, socket.SOCK_STREAM): 91 if socktype == socket.SOCK_STREAM: 92 af = family 93 addr = sockaddr 94 break 95 else: 96 raise exception.SSHError('No suitable address family for %s' % hostname) 97 sock = socket.socket(af, socket.SOCK_STREAM) 98 sock.settimeout(self._timeout) 99 sock.connect((hostname, port)) 100 return sock
101
102 - def _load_rsa_key(self, private_key, private_key_pass=None):
103 private_key_file = os.path.expanduser(private_key) 104 try: 105 rsa_key = paramiko.RSAKey.from_private_key_file(private_key_file, private_key_pass) 106 log.info("Using private key %s (rsa)" % private_key) 107 return rsa_key 108 except paramiko.SSHException,e: 109 log.error('invalid rsa key or password specified')
110
111 - def _load_dsa_key(self, private_key, private_key_pass=None):
112 private_key_file = os.path.expanduser(private_key) 113 try: 114 dsa_key = paramiko.DSSKey.from_private_key_file(private_key_file, private_key_pass) 115 log.info("Using private key %s (dsa)" % private_key) 116 return dsa_key 117 except paramiko.SSHException,e: 118 log.error('invalid dsa key or password specified')
119
120 - def _sftp_connect(self):
121 """Establish the SFTP connection.""" 122 if not self._sftp_live: 123 self._sftp = paramiko.SFTPClient.from_transport(self._transport) 124 self._sftp_live = True
125
126 - def remote_file(self, file, mode='w'):
127 """Returns a remote file descriptor""" 128 self._sftp_connect() 129 rfile = self._sftp.open(file, mode) 130 rfile.name=file 131 return rfile
132
133 - def path_exists(self, path):
134 """ 135 Test whether a remote path exists. Returns False for broken symbolic links 136 """ 137 self._sftp_connect() 138 try: 139 self.stat(path) 140 return True 141 except IOError,e: 142 return False
143
144 - def ls(self, path):
145 """ 146 Return a list containing the names of the entries in the remote path. 147 """ 148 self._sftp_connect() 149 return [os.path.join(path,f) for f in self._sftp.listdir(path)]
150
151 - def isdir(self, path):
152 """ 153 Return true if the remote path refers to an existing directory. 154 """ 155 self._sftp_connect() 156 try: 157 s = self.stat(path) 158 except IOError,e: 159 return False 160 return stat.S_ISDIR(s.st_mode)
161
162 - def isfile(self, path):
163 """ 164 Return true if the remote path refers to an existing file. 165 """ 166 self._sftp_connect() 167 try: 168 s = self.stat(path) 169 except IOError,e: 170 return False 171 return stat.S_ISREG(s.st_mode)
172
173 - def stat(self, path):
174 """ 175 Perform a stat system call on the given remote path. 176 """ 177 self._sftp_connect() 178 return self._sftp.stat(path)
179
180 - def get(self, remotepath, localpath = None):
181 """Copies a file between the remote host and the local host.""" 182 if not localpath: 183 localpath = os.path.split(remotepath)[1] 184 self._sftp_connect() 185 self._sftp.get(remotepath, localpath)
186
187 - def put(self, localpath, remotepath = None):
188 """Copies a file between the local host and the remote host.""" 189 if not remotepath: 190 remotepath = os.path.split(localpath)[1] 191 self._sftp_connect() 192 self._sftp.put(localpath, remotepath)
193
194 - def execute_async(self, command):
195 """ 196 Executes a remote command without blocking 197 198 NOTE: this method will not block, however, if your process does not 199 complete or background itself before the python process executing this 200 code exits, it will not persist on the remote machine 201 """ 202 203 channel = self._transport.open_session() 204 channel.exec_command(command)
205
206 - def execute(self, command, silent = True, only_printable = False, 207 ignore_exit_status=False):
208 """ 209 Execute a remote command and return stdout/stderr 210 211 NOTE: this function blocks until the process finishes 212 213 kwargs: 214 silent - do not print output 215 only_printable - filter the command's output to allow only printable 216 characters 217 returns List of output lines 218 """ 219 channel = self._transport.open_session() 220 channel.exec_command(command) 221 stdout = channel.makefile('rb', -1) 222 stderr = channel.makefile_stderr('rb', -1) 223 output = [] 224 line = None 225 if silent: 226 output = stdout.readlines() + stderr.readlines() 227 else: 228 while line != '': 229 line = stdout.readline() 230 if only_printable: 231 line = ''.join(char for char in line if char in string.printable) 232 if line != '': 233 output.append(line) 234 print line, 235 236 for line in stderr.readlines(): 237 output.append(line) 238 print line; 239 output = [ line.strip() for line in output ] 240 #TODO: warn about command failures 241 exit_status = channel.recv_exit_status() 242 if exit_status != 0: 243 if not ignore_exit_status: 244 log.error("command %s failed with status %d" % (command, 245 exit_status)) 246 else: 247 log.debug("command %s failed with status %d" % (command, 248 exit_status)) 249 if silent: 250 for line in output: 251 log.debug(line.strip()) 252 return output
253
254 - def close(self):
255 """Closes the connection and cleans up.""" 256 # Close SFTP Connection. 257 if self._sftp_live: 258 self._sftp.close() 259 self._sftp_live = False 260 # Close the SSH Transport. 261 if self._transport_live: 262 self._transport.close() 263 self._transport_live = False
264
265 - def interactive_shell(self):
266 try: 267 chan = self._transport.open_session() 268 chan.get_pty() 269 chan.invoke_shell() 270 log.info('Starting interactive shell...') 271 if HAS_TERMIOS: 272 self._posix_shell(chan) 273 else: 274 self._windows_shell(chan) 275 chan.close() 276 except Exception, e: 277 import traceback 278 print '*** Caught exception: %s: %s' % (e.__class__, e) 279 traceback.print_exc()
280
281 - def _posix_shell(self, chan):
282 import select 283 284 oldtty = termios.tcgetattr(sys.stdin) 285 try: 286 tty.setraw(sys.stdin.fileno()) 287 tty.setcbreak(sys.stdin.fileno()) 288 chan.settimeout(0.0) 289 290 chan.send('eval $(resize)\n') # needs to be sent to give vim correct size FIX 291 292 while True: 293 r, w, e = select.select([chan, sys.stdin], [], []) 294 if chan in r: 295 try: 296 x = chan.recv(1024) 297 if len(x) == 0: 298 print '\r\n*** EOF\r\n', 299 break 300 sys.stdout.write(x) 301 sys.stdout.flush() 302 except socket.timeout: 303 pass 304 if sys.stdin in r: 305 x = os.read(sys.stdin.fileno(),1) # fixes up arrow problem 306 if len(x) == 0: 307 break 308 chan.send(x) 309 finally: 310 termios.tcsetattr(sys.stdin, termios.TCSADRAIN, oldtty)
311 312 313 # thanks to Mike Looijmans for this code
314 - def _windows_shell(self, chan):
315 import threading 316 317 sys.stdout.write("Line-buffered terminal emulation. Press F6 or ^Z to send EOF.\r\n\r\n") 318 319 def writeall(sock): 320 while True: 321 data = sock.recv(256) 322 if not data: 323 sys.stdout.write('\r\n*** EOF ***\r\n\r\n') 324 sys.stdout.flush() 325 break 326 sys.stdout.write(data) 327 sys.stdout.flush()
328 329 writer = threading.Thread(target=writeall, args=(chan,)) 330 writer.start() 331 332 try: 333 while True: 334 d = sys.stdin.read(1) 335 if not d: 336 break 337 chan.send(d) 338 except EOFError: 339 # user hit ^Z or F6 340 pass
341
342 - def __del__(self):
343 """Attempt to clean up if not explicitly closed.""" 344 log.debug('__del__ called') 345 self.close()
346
347 -def main():
348 """Little test when called directly.""" 349 # Set these to your own details. 350 myssh = Connection('somehost.domain.com') 351 print myssh.execute('hostname') 352 #myssh.put('ssh.py') 353 myssh.close()
354 355 if __name__ == "__main__": 356 main() 357