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