1#!/usr/bin/env python
2
3r"""
4This module provides many valuable ssh functions such as sprint_connection,
5execute_ssh_command, etc.
6"""
7
8import sys
9import re
10import socket
11import paramiko
12import exceptions
13
14import gen_print as gp
15
16from robot.libraries.BuiltIn import BuiltIn
17from SSHLibrary import SSHLibrary
18sshlib = SSHLibrary()
19
20
21def sprint_connection(connection,
22                      indent=0):
23    r"""
24    sprint data from the connection object to a string and return it.
25
26    connection                      A connection object which is created by
27                                    the SSHlibrary open_connection() function.
28    indent                          The number of characters to indent the
29                                    output.
30    """
31
32    buffer = gp.sindent("", indent)
33    buffer += "connection:\n"
34    indent += 2
35    buffer += gp.sprint_varx("index", connection.index, 0, indent)
36    buffer += gp.sprint_varx("host", connection.host, 0, indent)
37    buffer += gp.sprint_varx("alias", connection.alias, 0, indent)
38    buffer += gp.sprint_varx("port", connection.port, 0, indent)
39    buffer += gp.sprint_varx("timeout", connection.timeout, 0, indent)
40    buffer += gp.sprint_varx("newline", connection.newline, 0, indent)
41    buffer += gp.sprint_varx("prompt", connection.prompt, 0, indent)
42    buffer += gp.sprint_varx("term_type", connection.term_type, 0, indent)
43    buffer += gp.sprint_varx("width", connection.width, 0, indent)
44    buffer += gp.sprint_varx("height", connection.height, 0, indent)
45    buffer += gp.sprint_varx("path_separator", connection.path_separator, 0,
46                             indent)
47    buffer += gp.sprint_varx("encoding", connection.encoding, 0, indent)
48
49    return buffer
50
51
52def sprint_connections(connections=None,
53                       indent=0):
54    r"""
55    sprint data from the connections list to a string and return it.
56
57    connections                     A list of connection objects which are
58                                    created by the SSHlibrary open_connection
59                                    function.  If this value is null, this
60                                    function will populate with a call to the
61                                    SSHlibrary get_connections() function.
62    indent                          The number of characters to indent the
63                                    output.
64    """
65
66    if connections is None:
67        connections = sshlib.get_connections()
68
69    buffer = ""
70    for connection in connections:
71        buffer += sprint_connection(connection, indent)
72
73    return buffer
74
75
76def find_connection(open_connection_args={}):
77    r"""
78    Find connection that matches the given connection arguments and return
79    connection object.  Return False if no matching connection is found.
80
81    Description of argument(s):
82    open_connection_args            A dictionary of arg names and values which
83                                    are legal to pass to the SSHLibrary
84                                    open_connection function as parms/args.
85                                    For a match to occur, the value for each
86                                    item in open_connection_args must match
87                                    the corresponding value in the connection
88                                    being examined.
89    """
90
91    global sshlib
92
93    for connection in sshlib.get_connections():
94        # Create connection_dict from connection object.
95        connection_dict = dict((key, str(value)) for key, value in
96                               connection._config.iteritems())
97        if dict(connection_dict, **open_connection_args) == connection_dict:
98            return connection
99
100    return False
101
102
103def login_ssh(login_args={},
104              max_login_attempts=5):
105    r"""
106    Login on the latest open SSH connection.  Retry on failure up to
107    max_login_attempts.
108
109    The caller is responsible for making sure there is an open SSH connection.
110
111    Description of argument(s):
112    login_args                      A dictionary containing the key/value
113                                    pairs which are acceptable to the
114                                    SSHLibrary login function as parms/args.
115                                    At a minimum, this should contain a
116                                    'username' and a 'password' entry.
117    max_login_attempts              The max number of times to try logging in
118                                    (in the event of login failures).
119    """
120
121    gp.lprint_executing()
122
123    global sshlib
124
125    # Get connection data for debug output.
126    connection = sshlib.get_connection()
127    gp.lprintn(sprint_connection(connection))
128    for login_attempt_num in range(1, max_login_attempts + 1):
129        gp.lprint_timen("Logging in to " + connection.host + ".")
130        gp.lprint_var(login_attempt_num)
131        try:
132            out_buf = sshlib.login(**login_args)
133        except Exception as login_exception:
134            # Login will sometimes fail if the connection is new.
135            except_type, except_value, except_traceback = sys.exc_info()
136            gp.lprint_var(except_type)
137            gp.lprint_varx("except_value", str(except_value))
138            if except_type is paramiko.ssh_exception.SSHException and\
139                    re.match(r"No existing session", str(except_value)):
140                continue
141            else:
142                # We don't tolerate any other error so break from loop and
143                # re-raise exception.
144                break
145        # If we get to this point, the login has worked and we can return.
146        gp.lpvar(out_buf)
147        return
148
149    # If we get to this point, the login has failed on all attempts so the
150    # exception will be raised again.
151    raise(login_exception)
152
153
154def execute_ssh_command(cmd_buf,
155                        open_connection_args={},
156                        login_args={},
157                        print_out=0,
158                        print_err=0,
159                        ignore_err=1,
160                        fork=0,
161                        quiet=None,
162                        test_mode=None):
163    r"""
164    Run the given command in an SSH session and return the stdout, stderr and
165    the return code.
166
167    If there is no open SSH connection, this function will connect and login.
168    Likewise, if the caller has not yet logged in to the connection, this
169    function will do the login.
170
171    Description of arguments:
172    cmd_buf                         The command string to be run in an SSH
173                                    session.
174    open_connection_args            A dictionary of arg names and values which
175                                    are legal to pass to the SSHLibrary
176                                    open_connection function as parms/args.
177                                    At a minimum, this should contain a 'host'
178                                    entry.
179    login_args                      A dictionary containing the key/value
180                                    pairs which are acceptable to the
181                                    SSHLibrary login function as parms/args.
182                                    At a minimum, this should contain a
183                                    'username' and a 'password' entry.
184    print_out                       If this is set, this function will print
185                                    the stdout/stderr generated by the shell
186                                    command.
187    print_err                       If show_err is set, this function will
188                                    print a standardized error report if the
189                                    shell command returns non-zero.
190    ignore_err                      Indicates that errors encountered on the
191                                    sshlib.execute_command are to be ignored.
192    fork                            Indicates that sshlib.start is to be used
193                                    rather than sshlib.execute_command.
194    quiet                           Indicates whether this function should run
195                                    the pissuing() function which prints an
196                                    "Issuing: <cmd string>" to stdout.  This
197                                    defaults to the global quiet value.
198    test_mode                       If test_mode is set, this function will
199                                    not actually run the command.  This
200                                    defaults to the global test_mode value.
201    """
202
203    gp.lprint_executing()
204
205    # Obtain default values.
206    quiet = int(gp.get_var_value(quiet, 0))
207    test_mode = int(gp.get_var_value(test_mode, 0))
208
209    if not quiet:
210        gp.pissuing(cmd_buf, test_mode)
211    gp.lpissuing(cmd_buf, test_mode)
212
213    if test_mode:
214        return "", "", 0
215
216    global sshlib
217
218    # Look for existing SSH connection.
219    # Prepare a search connection dictionary.
220    search_connection_args = open_connection_args.copy()
221    # Remove keys that don't work well for searches.
222    search_connection_args.pop("timeout", None)
223    connection = find_connection(search_connection_args)
224    if connection:
225        gp.lprint_timen("Found the following existing connection:")
226        gp.lprintn(sprint_connection(connection))
227        if connection.alias == "":
228            index_or_alias = connection.index
229        else:
230            index_or_alias = connection.alias
231        gp.lprint_timen("Switching to existing connection: \"" +
232                        str(index_or_alias) + "\".")
233        sshlib.switch_connection(index_or_alias)
234    else:
235        gp.lprint_timen("Connecting to " + open_connection_args['host'] + ".")
236        cix = sshlib.open_connection(**open_connection_args)
237        login_ssh(login_args)
238
239    max_exec_cmd_attempts = 2
240    for exec_cmd_attempt_num in range(1, max_exec_cmd_attempts + 1):
241        gp.lprint_var(exec_cmd_attempt_num)
242        try:
243            if fork:
244                sshlib.start_command(cmd_buf)
245            else:
246                stdout, stderr, rc = sshlib.execute_command(cmd_buf,
247                                                            return_stdout=True,
248                                                            return_stderr=True,
249                                                            return_rc=True)
250        except Exception as execute_exception:
251            except_type, except_value, except_traceback = sys.exc_info()
252            gp.lprint_var(except_type)
253            gp.lprint_varx("except_value", str(except_value))
254
255            if except_type is exceptions.AssertionError and\
256               re.match(r"Connection not open", str(except_value)):
257                login_ssh(login_args)
258                # Now we must continue to next loop iteration to retry the
259                # execute_command.
260                continue
261            if (except_type is paramiko.ssh_exception.SSHException and
262                re.match(r"SSH session not active", str(except_value))) or\
263               (except_type is socket.error and
264                re.match(r"\[Errno 104\] Connection reset by peer",
265                         str(except_value))):
266                # Close and re-open a connection.
267                # Note: close_connection() doesn't appear to get rid of the
268                # connection.  It merely closes it.  Since there is a concern
269                # about over-consumption of resources, we use
270                # close_all_connections() which also gets rid of all
271                # connections.
272                gp.lprint_timen("Closing all connections.")
273                sshlib.close_all_connections()
274                gp.lprint_timen("Connecting to " +
275                                open_connection_args['host'] + ".")
276                cix = sshlib.open_connection(**open_connection_args)
277                login_ssh(login_args)
278                continue
279
280            # We do not handle any other RuntimeErrors so we will raise the
281            # exception again.
282            sshlib.close_all_connections()
283            raise(execute_exception)
284
285        # If we get to this point, the command was executed.
286        break
287
288    if fork:
289        return
290
291    if rc != 0 and print_err:
292        gp.print_var(rc, 1)
293        if not print_out:
294            gp.print_var(stderr)
295            gp.print_var(stdout)
296
297    if print_out:
298        gp.printn(stderr + stdout)
299
300    if not ignore_err:
301        message = gp.sprint_error("The prior SSH" +
302                                  " command returned a non-zero return" +
303                                  " code:\n" + gp.sprint_var(rc, 1) + stderr +
304                                  "\n")
305        BuiltIn().should_be_equal(rc, 0, message)
306
307    return stdout, stderr, rc
308