/boto-2.5.2/boto/manage/cmdshell.py

  1# Copyright (c) 2006-2009 Mitch Garnaat http://garnaat.org/
  2#
  3# Permission is hereby granted, free of charge, to any person obtaining a
  4# copy of this software and associated documentation files (the
  5# "Software"), to deal in the Software without restriction, including
  6# without limitation the rights to use, copy, modify, merge, publish, dis-
  7# tribute, sublicense, and/or sell copies of the Software, and to permit
  8# persons to whom the Software is furnished to do so, subject to the fol-
  9# lowing conditions:
 10#
 11# The above copyright notice and this permission notice shall be included
 12# in all copies or substantial portions of the Software.
 13#
 14# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
 15# OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABIL-
 16# ITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT
 17# SHALL THE AUTHOR BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, 
 18# WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 19# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
 20# IN THE SOFTWARE.
 21
 22from boto.mashups.interactive import interactive_shell
 23import boto
 24import os
 25import time
 26import shutil
 27import StringIO
 28import paramiko
 29import socket
 30import subprocess
 31
 32
 33class SSHClient(object):
 34
 35    def __init__(self, server,
 36                 host_key_file='~/.ssh/known_hosts',
 37                 uname='root', ssh_pwd=None):
 38        self.server = server
 39        self.host_key_file = host_key_file
 40        self.uname = uname
 41        self._pkey = paramiko.RSAKey.from_private_key_file(server.ssh_key_file,
 42                                                           password=ssh_pwd)
 43        self._ssh_client = paramiko.SSHClient()
 44        self._ssh_client.load_system_host_keys()
 45        self._ssh_client.load_host_keys(os.path.expanduser(host_key_file))
 46        self._ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
 47        self.connect()
 48
 49    def connect(self):
 50        retry = 0
 51        while retry < 5:
 52            try:
 53                self._ssh_client.connect(self.server.hostname,
 54                                         username=self.uname,
 55                                         pkey=self._pkey)
 56                return
 57            except socket.error, (value, message):
 58                if value == 61 or value == 111:
 59                    print 'SSH Connection refused, will retry in 5 seconds'
 60                    time.sleep(5)
 61                    retry += 1
 62                else:
 63                    raise
 64            except paramiko.BadHostKeyException:
 65                print "%s has an entry in ~/.ssh/known_hosts and it doesn't match" % self.server.hostname
 66                print 'Edit that file to remove the entry and then hit return to try again'
 67                raw_input('Hit Enter when ready')
 68                retry += 1
 69            except EOFError:
 70                print 'Unexpected Error from SSH Connection, retry in 5 seconds'
 71                time.sleep(5)
 72                retry += 1
 73        print 'Could not establish SSH connection'
 74
 75    def open_sftp(self):
 76        return self._ssh_client.open_sftp()
 77
 78    def get_file(self, src, dst):
 79        sftp_client = self.open_sftp()
 80        sftp_client.get(src, dst)
 81
 82    def put_file(self, src, dst):
 83        sftp_client = self.open_sftp()
 84        sftp_client.put(src, dst)
 85
 86    def open(self, filename, mode='r', bufsize=-1):
 87        """
 88        Open a file on the remote system and return a file-like object.
 89        """
 90        sftp_client = self.open_sftp()
 91        return sftp_client.open(filename, mode, bufsize)
 92
 93    def listdir(self, path):
 94        sftp_client = self.open_sftp()
 95        return sftp_client.listdir(path)
 96
 97    def isdir(self, path):
 98        status = self.run('[ -d %s ] || echo "FALSE"' % path)
 99        if status[1].startswith('FALSE'):
100            return 0
101        return 1
102
103    def exists(self, path):
104        status = self.run('[ -a %s ] || echo "FALSE"' % path)
105        if status[1].startswith('FALSE'):
106            return 0
107        return 1
108
109    def shell(self):
110        """
111        Start an interactive shell session on the remote host.
112        """
113        channel = self._ssh_client.invoke_shell()
114        interactive_shell(channel)
115
116    def run(self, command):
117        """
118        Execute a command on the remote host.  Return a tuple containing
119        an integer status and a two strings, the first containing stdout
120        and the second containing stderr from the command.
121        """
122        boto.log.debug('running:%s on %s' % (command, self.server.instance_id))
123        status = 0
124        try:
125            t = self._ssh_client.exec_command(command)
126        except paramiko.SSHException:
127            status = 1
128        std_out = t[1].read()
129        std_err = t[2].read()
130        t[0].close()
131        t[1].close()
132        t[2].close()
133        boto.log.debug('stdout: %s' % std_out)
134        boto.log.debug('stderr: %s' % std_err)
135        return (status, std_out, std_err)
136
137    def run_pty(self, command):
138        """
139        Execute a command on the remote host with a pseudo-terminal.
140        Returns a string containing the output of the command.
141        """
142        boto.log.debug('running:%s on %s' % (command, self.server.instance_id))
143        channel = self._ssh_client.get_transport().open_session()
144        channel.get_pty()
145        channel.exec_command(command)
146        return channel
147
148    def close(self):
149        transport = self._ssh_client.get_transport()
150        transport.close()
151        self.server.reset_cmdshell()
152
153class LocalClient(object):
154
155    def __init__(self, server, host_key_file=None, uname='root'):
156        self.server = server
157        self.host_key_file = host_key_file
158        self.uname = uname
159
160    def get_file(self, src, dst):
161        shutil.copyfile(src, dst)
162
163    def put_file(self, src, dst):
164        shutil.copyfile(src, dst)
165
166    def listdir(self, path):
167        return os.listdir(path)
168
169    def isdir(self, path):
170        return os.path.isdir(path)
171
172    def exists(self, path):
173        return os.path.exists(path)
174
175    def shell(self):
176        raise NotImplementedError('shell not supported with LocalClient')
177
178    def run(self):
179        boto.log.info('running:%s' % self.command)
180        log_fp = StringIO.StringIO()
181        process = subprocess.Popen(self.command, shell=True, stdin=subprocess.PIPE,
182                                   stdout=subprocess.PIPE, stderr=subprocess.PIPE)
183        while process.poll() == None:
184            time.sleep(1)
185            t = process.communicate()
186            log_fp.write(t[0])
187            log_fp.write(t[1])
188        boto.log.info(log_fp.getvalue())
189        boto.log.info('output: %s' % log_fp.getvalue())
190        return (process.returncode, log_fp.getvalue())
191
192    def close(self):
193        pass
194
195class FakeServer(object):
196    """
197    A little class to fake out SSHClient (which is expecting a
198    :class`boto.manage.server.Server` instance.  This allows us
199    to 
200    """
201    def __init__(self, instance, ssh_key_file):
202        self.instance = instance
203        self.ssh_key_file = ssh_key_file
204        self.hostname = instance.dns_name
205        self.instance_id = self.instance.id
206        
207def start(server):
208    instance_id = boto.config.get('Instance', 'instance-id', None)
209    if instance_id == server.instance_id:
210        return LocalClient(server)
211    else:
212        return SSHClient(server)
213
214def sshclient_from_instance(instance, ssh_key_file,
215                            host_key_file='~/.ssh/known_hosts',
216                            user_name='root', ssh_pwd=None):
217    """
218    Create and return an SSHClient object given an
219    instance object.
220
221    :type instance: :class`boto.ec2.instance.Instance` object
222    :param instance: The instance object.
223
224    :type ssh_key_file: str
225    :param ssh_key_file: A path to the private key file used
226                         to log into instance.
227
228    :type host_key_file: str
229    :param host_key_file: A path to the known_hosts file used
230                          by the SSH client.
231                          Defaults to ~/.ssh/known_hosts
232    :type user_name: str
233    :param user_name: The username to use when logging into
234                      the instance.  Defaults to root.
235
236    :type ssh_pwd: str
237    :param ssh_pwd: The passphrase, if any, associated with
238                    private key.
239    """
240    s = FakeServer(instance, ssh_key_file)
241    return SSHClient(s, host_key_file, user_name, ssh_pwd)