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