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    NOTE: There is special handling when open_connection_args['alias'] equals
172    "device_connection".
173    - A write, rather than an execute_command, is done.
174    - Only stdout is returned (no stderr or rc).
175    - print_err, ignore_err and fork are not supported.
176
177    Description of arguments:
178    cmd_buf                         The command string to be run in an SSH
179                                    session.
180    open_connection_args            A dictionary of arg names and values which
181                                    are legal to pass to the SSHLibrary
182                                    open_connection function as parms/args.
183                                    At a minimum, this should contain a 'host'
184                                    entry.
185    login_args                      A dictionary containing the key/value
186                                    pairs which are acceptable to the
187                                    SSHLibrary login function as parms/args.
188                                    At a minimum, this should contain a
189                                    'username' and a 'password' entry.
190    print_out                       If this is set, this function will print
191                                    the stdout/stderr generated by the shell
192                                    command.
193    print_err                       If show_err is set, this function will
194                                    print a standardized error report if the
195                                    shell command returns non-zero.
196    ignore_err                      Indicates that errors encountered on the
197                                    sshlib.execute_command are to be ignored.
198    fork                            Indicates that sshlib.start is to be used
199                                    rather than sshlib.execute_command.
200    quiet                           Indicates whether this function should run
201                                    the pissuing() function which prints an
202                                    "Issuing: <cmd string>" to stdout.  This
203                                    defaults to the global quiet value.
204    test_mode                       If test_mode is set, this function will
205                                    not actually run the command.  This
206                                    defaults to the global test_mode value.
207    """
208
209    gp.lprint_executing()
210
211    # Obtain default values.
212    quiet = int(gp.get_var_value(quiet, 0))
213    test_mode = int(gp.get_var_value(test_mode, 0))
214
215    if not quiet:
216        gp.pissuing(cmd_buf, test_mode)
217    gp.lpissuing(cmd_buf, test_mode)
218
219    if test_mode:
220        return "", "", 0
221
222    global sshlib
223
224    # Look for existing SSH connection.
225    # Prepare a search connection dictionary.
226    search_connection_args = open_connection_args.copy()
227    # Remove keys that don't work well for searches.
228    search_connection_args.pop("timeout", None)
229    connection = find_connection(search_connection_args)
230    if connection:
231        gp.lprint_timen("Found the following existing connection:")
232        gp.lprintn(sprint_connection(connection))
233        if connection.alias == "":
234            index_or_alias = connection.index
235        else:
236            index_or_alias = connection.alias
237        gp.lprint_timen("Switching to existing connection: \""
238                        + str(index_or_alias) + "\".")
239        sshlib.switch_connection(index_or_alias)
240    else:
241        gp.lprint_timen("Connecting to " + open_connection_args['host'] + ".")
242        cix = sshlib.open_connection(**open_connection_args)
243        login_ssh(login_args)
244
245    max_exec_cmd_attempts = 2
246    for exec_cmd_attempt_num in range(1, max_exec_cmd_attempts + 1):
247        gp.lprint_var(exec_cmd_attempt_num)
248        try:
249            if fork:
250                sshlib.start_command(cmd_buf)
251            else:
252                if open_connection_args['alias'] == "device_connection":
253                    stdout = sshlib.write(cmd_buf)
254                    stderr = ""
255                    rc = 0
256                else:
257                    stdout, stderr, rc = \
258                        sshlib.execute_command(cmd_buf,
259                                               return_stdout=True,
260                                               return_stderr=True,
261                                               return_rc=True)
262        except Exception as execute_exception:
263            except_type, except_value, except_traceback = sys.exc_info()
264            gp.lprint_var(except_type)
265            gp.lprint_varx("except_value", str(except_value))
266
267            if except_type is exceptions.AssertionError and\
268               re.match(r"Connection not open", str(except_value)):
269                login_ssh(login_args)
270                # Now we must continue to next loop iteration to retry the
271                # execute_command.
272                continue
273            if (except_type is paramiko.ssh_exception.SSHException
274                and re.match(r"SSH session not active", str(except_value))) or\
275               (except_type is socket.error
276                and re.match(r"\[Errno 104\] Connection reset by peer",
277                             str(except_value))):
278                # Close and re-open a connection.
279                # Note: close_connection() doesn't appear to get rid of the
280                # connection.  It merely closes it.  Since there is a concern
281                # about over-consumption of resources, we use
282                # close_all_connections() which also gets rid of all
283                # connections.
284                gp.lprint_timen("Closing all connections.")
285                sshlib.close_all_connections()
286                gp.lprint_timen("Connecting to "
287                                + open_connection_args['host'] + ".")
288                cix = sshlib.open_connection(**open_connection_args)
289                login_ssh(login_args)
290                continue
291
292            # We do not handle any other RuntimeErrors so we will raise the
293            # exception again.
294            sshlib.close_all_connections()
295            raise(execute_exception)
296
297        # If we get to this point, the command was executed.
298        break
299
300    if fork:
301        return
302
303    if rc != 0 and print_err:
304        gp.print_var(rc, 1)
305        if not print_out:
306            gp.print_var(stderr)
307            gp.print_var(stdout)
308
309    if print_out:
310        gp.printn(stderr + stdout)
311
312    if not ignore_err:
313        message = gp.sprint_error("The prior SSH"
314                                  + " command returned a non-zero return"
315                                  + " code:\n" + gp.sprint_var(rc, 1) + stderr
316                                  + "\n")
317        BuiltIn().should_be_equal(rc, 0, message)
318
319    if open_connection_args['alias'] == "device_connection":
320        return stdout
321    return stdout, stderr, rc
322