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"""
22The cmdshell module uses the paramiko package to create SSH connections
23to the servers that are represented by instance objects. The module has
24functions for running commands, managing files, and opening interactive
25shell sessions over those connections.
26"""
27from boto.mashups.interactive import interactive_shell
28import boto
29import os
30import time
31import shutil
32import paramiko
33import socket
34import subprocess
35
36from boto.compat import StringIO
37
38class SSHClient(object):
39    """
40    This class creates a paramiko.SSHClient() object that represents
41    a session with an SSH server. You can use the SSHClient object to send
42    commands to the remote host and manipulate files on the remote host.
43
44    :ivar server: A Server object or FakeServer object.
45    :ivar host_key_file: The path to the user's .ssh key files.
46    :ivar uname: The username for the SSH connection. Default = 'root'.
47    :ivar timeout: The optional timeout variable for the TCP connection.
48    :ivar ssh_pwd: An optional password to use for authentication or for
49                    unlocking the private key.
50    """
51    def __init__(self, server,
52                 host_key_file='~/.ssh/known_hosts',
53                 uname='root', timeout=None, ssh_pwd=None):
54        self.server = server
55        self.host_key_file = host_key_file
56        self.uname = uname
57        self._timeout = timeout
58        self._pkey = paramiko.RSAKey.from_private_key_file(server.ssh_key_file,
59                                                           password=ssh_pwd)
60        self._ssh_client = paramiko.SSHClient()
61        self._ssh_client.load_system_host_keys()
62        self._ssh_client.load_host_keys(os.path.expanduser(host_key_file))
63        self._ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
64        self.connect()
65
66    def connect(self, num_retries=5):
67        """
68        Connect to an SSH server and authenticate with it.
69
70        :type num_retries: int
71        :param num_retries: The maximum number of connection attempts.
72        """
73        retry = 0
74        while retry < num_retries:
75            try:
76                self._ssh_client.connect(self.server.hostname,
77                                         username=self.uname,
78                                         pkey=self._pkey,
79                                         timeout=self._timeout)
80                return
81            except socket.error as xxx_todo_changeme:
82                (value, message) = xxx_todo_changeme.args
83                if value in (51, 61, 111):
84                    print('SSH Connection refused, will retry in 5 seconds')
85                    time.sleep(5)
86                    retry += 1
87                else:
88                    raise
89            except paramiko.BadHostKeyException:
90                print("%s has an entry in ~/.ssh/known_hosts and it doesn't match" % self.server.hostname)
91                print('Edit that file to remove the entry and then hit return to try again')
92                raw_input('Hit Enter when ready')
93                retry += 1
94            except EOFError:
95                print('Unexpected Error from SSH Connection, retry in 5 seconds')
96                time.sleep(5)
97                retry += 1
98        print('Could not establish SSH connection')
99
100    def open_sftp(self):
101        """
102        Open an SFTP session on the SSH server.
103
104        :rtype: :class:`paramiko.sftp_client.SFTPClient`
105        :return: An SFTP client object.
106        """
107        return self._ssh_client.open_sftp()
108
109    def get_file(self, src, dst):
110        """
111        Open an SFTP session on the remote host, and copy a file from
112        the remote host to the specified path on the local host.
113
114        :type src: string
115        :param src: The path to the target file on the remote host.
116
117        :type dst: string
118        :param dst: The path on your local host where you want to
119                    store the file.
120        """
121        sftp_client = self.open_sftp()
122        sftp_client.get(src, dst)
123
124    def put_file(self, src, dst):
125        """
126        Open an SFTP session on the remote host, and copy a file from
127        the local host to the specified path on the remote host.
128
129        :type src: string
130        :param src: The path to the target file on your local host.
131
132        :type dst: string
133        :param dst: The path on the remote host where you want to store
134                    the file.
135        """
136        sftp_client = self.open_sftp()
137        sftp_client.put(src, dst)
138
139    def open(self, filename, mode='r', bufsize=-1):
140        """
141        Open an SFTP session to the remote host, and open a file on
142        that host.
143
144        :type filename: string
145        :param filename: The path to the file on the remote host.
146
147        :type mode: string
148        :param mode: The file interaction mode.
149
150        :type bufsize: integer
151        :param bufsize: The file buffer size.
152
153        :rtype: :class:`paramiko.sftp_file.SFTPFile`
154        :return: A paramiko proxy object for a file on the remote server.
155        """
156        sftp_client = self.open_sftp()
157        return sftp_client.open(filename, mode, bufsize)
158
159    def listdir(self, path):
160        """
161        List all of the files and subdirectories at the specified path
162        on the remote host.
163
164        :type path: string
165        :param path: The base path from which to obtain the list.
166
167        :rtype: list
168        :return: A list of files and subdirectories at the specified path.
169        """
170        sftp_client = self.open_sftp()
171        return sftp_client.listdir(path)
172
173    def isdir(self, path):
174        """
175        Check the specified path on the remote host to determine if
176        it is a directory.
177
178        :type path: string
179        :param path: The path to the directory that you want to check.
180
181        :rtype: integer
182        :return: If the path is a directory, the function returns 1.
183                If the path is a file or an invalid path, the function
184                returns 0.
185        """
186        status = self.run('[ -d %s ] || echo "FALSE"' % path)
187        if status[1].startswith('FALSE'):
188            return 0
189        return 1
190
191    def exists(self, path):
192        """
193        Check the remote host for the specified path, or a file
194        at the specified path. This function returns 1 if the
195        path or the file exist on the remote host, and returns 0 if
196        the path or the file does not exist on the remote host.
197
198        :type path: string
199        :param path: The path to the directory or file that you want to check.
200
201        :rtype: integer
202        :return: If the path or the file exist, the function returns 1.
203                If the path or the file do not exist on the remote host,
204                the function returns 0.
205        """
206
207        status = self.run('[ -a %s ] || echo "FALSE"' % path)
208        if status[1].startswith('FALSE'):
209            return 0
210        return 1
211
212    def shell(self):
213        """
214        Start an interactive shell session with the remote host.
215        """
216        channel = self._ssh_client.invoke_shell()
217        interactive_shell(channel)
218
219    def run(self, command):
220        """
221        Run a command on the remote host.
222
223        :type command: string
224        :param command: The command that you want to send to the remote host.
225
226        :rtype: tuple
227        :return: This function returns a tuple that contains an integer status,
228                the stdout from the command, and the stderr from the command.
229
230        """
231        boto.log.debug('running:%s on %s' % (command, self.server.instance_id))
232        status = 0
233        try:
234            t = self._ssh_client.exec_command(command)
235        except paramiko.SSHException:
236            status = 1
237        std_out = t[1].read()
238        std_err = t[2].read()
239        t[0].close()
240        t[1].close()
241        t[2].close()
242        boto.log.debug('stdout: %s' % std_out)
243        boto.log.debug('stderr: %s' % std_err)
244        return (status, std_out, std_err)
245
246    def run_pty(self, command):
247        """
248        Request a pseudo-terminal from a server, and execute a command on that
249        server.
250
251        :type command: string
252        :param command: The command that you want to run on the remote host.
253
254        :rtype: :class:`paramiko.channel.Channel`
255        :return: An open channel object.
256        """
257        boto.log.debug('running:%s on %s' % (command, self.server.instance_id))
258        channel = self._ssh_client.get_transport().open_session()
259        channel.get_pty()
260        channel.exec_command(command)
261        return channel
262
263    def close(self):
264        """
265        Close an SSH session and any open channels that are tied to it.
266        """
267        transport = self._ssh_client.get_transport()
268        transport.close()
269        self.server.reset_cmdshell()
270
271class LocalClient(object):
272    """
273    :ivar server: A Server object or FakeServer object.
274    :ivar host_key_file: The path to the user's .ssh key files.
275    :ivar uname: The username for the SSH connection. Default = 'root'.
276    """
277    def __init__(self, server, host_key_file=None, uname='root'):
278        self.server = server
279        self.host_key_file = host_key_file
280        self.uname = uname
281
282    def get_file(self, src, dst):
283        """
284        Copy a file from one directory to another.
285        """
286        shutil.copyfile(src, dst)
287
288    def put_file(self, src, dst):
289        """
290        Copy a file from one directory to another.
291        """
292        shutil.copyfile(src, dst)
293
294    def listdir(self, path):
295        """
296        List all of the files and subdirectories at the specified path.
297
298        :rtype: list
299        :return: Return a list containing the names of the entries
300                in the directory given by path.
301        """
302        return os.listdir(path)
303
304    def isdir(self, path):
305        """
306        Check the specified path to determine if it is a directory.
307
308        :rtype: boolean
309        :return: Returns True if the path is an existing directory.
310        """
311        return os.path.isdir(path)
312
313    def exists(self, path):
314        """
315        Check for the specified path, or check a file at the specified path.
316
317        :rtype: boolean
318        :return: If the path or the file exist, the function returns True.
319        """
320        return os.path.exists(path)
321
322    def shell(self):
323        raise NotImplementedError('shell not supported with LocalClient')
324
325    def run(self):
326        """
327        Open a subprocess and run a command on the local host.
328
329        :rtype: tuple
330        :return: This function returns a tuple that contains an integer status
331                and a string with the combined stdout and stderr output.
332        """
333        boto.log.info('running:%s' % self.command)
334        log_fp = StringIO()
335        process = subprocess.Popen(self.command, shell=True, stdin=subprocess.PIPE,
336                                   stdout=subprocess.PIPE, stderr=subprocess.PIPE)
337        while process.poll() is None:
338            time.sleep(1)
339            t = process.communicate()
340            log_fp.write(t[0])
341            log_fp.write(t[1])
342        boto.log.info(log_fp.getvalue())
343        boto.log.info('output: %s' % log_fp.getvalue())
344        return (process.returncode, log_fp.getvalue())
345
346    def close(self):
347        pass
348
349class FakeServer(object):
350    """
351    This object has a subset of the variables that are normally in a
352    :class:`boto.manage.server.Server` object. You can use this FakeServer
353    object to create a :class:`boto.manage.SSHClient` object if you
354    don't have a real Server object.
355
356    :ivar instance: A boto Instance object.
357    :ivar ssh_key_file: The path to the SSH key file.
358    """
359    def __init__(self, instance, ssh_key_file):
360        self.instance = instance
361        self.ssh_key_file = ssh_key_file
362        self.hostname = instance.dns_name
363        self.instance_id = self.instance.id
364
365def start(server):
366    """
367    Connect to the specified server.
368
369    :return: If the server is local, the function returns a
370            :class:`boto.manage.cmdshell.LocalClient` object.
371            If the server is remote, the function returns a
372            :class:`boto.manage.cmdshell.SSHClient` object.
373    """
374    instance_id = boto.config.get('Instance', 'instance-id', None)
375    if instance_id == server.instance_id:
376        return LocalClient(server)
377    else:
378        return SSHClient(server)
379
380def sshclient_from_instance(instance, ssh_key_file,
381                            host_key_file='~/.ssh/known_hosts',
382                            user_name='root', ssh_pwd=None):
383    """
384    Create and return an SSHClient object given an
385    instance object.
386
387    :type instance: :class`boto.ec2.instance.Instance` object
388    :param instance: The instance object.
389
390    :type ssh_key_file: string
391    :param ssh_key_file: A path to the private key file that is
392                        used to log into the instance.
393
394    :type host_key_file: string
395    :param host_key_file: A path to the known_hosts file used
396                          by the SSH client.
397                          Defaults to ~/.ssh/known_hosts
398    :type user_name: string
399    :param user_name: The username to use when logging into
400                      the instance.  Defaults to root.
401
402    :type ssh_pwd: string
403    :param ssh_pwd: The passphrase, if any, associated with
404                    private key.
405    """
406    s = FakeServer(instance, ssh_key_file)
407    return SSHClient(s, host_key_file, user_name, ssh_pwd)
408