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