1#!/usr/bin/env python3
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            BuiltIn().log_to_console(out_buf)
129        except Exception:
130            # Login will sometimes fail if the connection is new.
131            except_type, except_value, except_traceback = sys.exc_info()
132            gp.lprint_var(except_type)
133            gp.lprint_varx("except_value", str(except_value))
134            if except_type is paramiko.ssh_exception.SSHException and\
135                    re.match(r"No existing session", str(except_value)):
136                continue
137            else:
138                # We don't tolerate any other error so break from loop and re-raise exception.
139                break
140        # If we get to this point, the login has worked and we can return.
141        gp.lprint_var(out_buf)
142        return
143
144    # If we get to this point, the login has failed on all attempts so the exception will be raised again.
145    raise(except_value)
146
147
148def execute_ssh_command(cmd_buf,
149                        open_connection_args={},
150                        login_args={},
151                        print_out=0,
152                        print_err=0,
153                        ignore_err=1,
154                        fork=0,
155                        quiet=None,
156                        test_mode=None,
157                        time_out=None):
158    r"""
159    Run the given command in an SSH session and return the stdout, stderr and the return code.
160
161    If there is no open SSH connection, this function will connect and login.  Likewise, if the caller has
162    not yet logged in to the connection, this function will do the login.
163
164    NOTE: There is special handling when open_connection_args['alias'] equals "device_connection".
165    - A write, rather than an execute_command, is done.
166    - Only stdout is returned (no stderr or rc).
167    - print_err, ignore_err and fork are not supported.
168
169    Description of arguments:
170    cmd_buf                         The command string to be run in an SSH session.
171    open_connection_args            A dictionary of arg names and values which are legal to pass to the
172                                    SSHLibrary open_connection function as parms/args.  At a minimum, this
173                                    should contain a 'host' entry.
174    login_args                      A dictionary containing the key/value pairs which are acceptable to the
175                                    SSHLibrary login function as parms/args.  At a minimum, this should
176                                    contain a 'username' and a 'password' entry.
177    print_out                       If this is set, this function will print the stdout/stderr generated by
178                                    the shell command.
179    print_err                       If show_err is set, this function will print a standardized error report
180                                    if the shell command returns non-zero.
181    ignore_err                      Indicates that errors encountered on the sshlib.execute_command are to be
182                                    ignored.
183    fork                            Indicates that sshlib.start is to be used rather than
184                                    sshlib.execute_command.
185    quiet                           Indicates whether this function should run the pissuing() function which
186                                    prints an "Issuing: <cmd string>" to stdout.  This defaults to the global
187                                    quiet value.
188    test_mode                       If test_mode is set, this function will not actually run the command.
189                                    This defaults to the global test_mode value.
190    time_out                        The amount of time to allow for the execution of cmd_buf.  A value of
191                                    None means that there is no limit to how long the command may take.
192    """
193
194    gp.lprint_executing()
195
196    # Obtain default values.
197    quiet = int(gp.get_var_value(quiet, 0))
198    test_mode = int(gp.get_var_value(test_mode, 0))
199
200    if not quiet:
201        gp.pissuing(cmd_buf, test_mode)
202    gp.lpissuing(cmd_buf, test_mode)
203
204    if test_mode:
205        return "", "", 0
206
207    global sshlib
208
209    max_exec_cmd_attempts = 2
210    # Look for existing SSH connection.
211    # Prepare a search connection dictionary.
212    search_connection_args = open_connection_args.copy()
213    # Remove keys that don't work well for searches.
214    search_connection_args.pop("timeout", None)
215    connection = find_connection(search_connection_args)
216    if connection:
217        gp.lprint_timen("Found the following existing connection:")
218        gp.lprintn(sprint_connection(connection))
219        if connection.alias == "":
220            index_or_alias = connection.index
221        else:
222            index_or_alias = connection.alias
223        gp.lprint_timen("Switching to existing connection: \""
224                        + str(index_or_alias) + "\".")
225        sshlib.switch_connection(index_or_alias)
226    else:
227        gp.lprint_timen("Connecting to " + open_connection_args['host'] + ".")
228        cix = sshlib.open_connection(**open_connection_args)
229        try:
230            login_ssh(login_args)
231        except Exception:
232            except_type, except_value, except_traceback = sys.exc_info()
233            rc = 1
234            stderr = str(except_value)
235            stdout = ""
236            max_exec_cmd_attempts = 0
237
238    for exec_cmd_attempt_num in range(1, max_exec_cmd_attempts + 1):
239        gp.lprint_var(exec_cmd_attempt_num)
240        try:
241            if fork:
242                sshlib.start_command(cmd_buf)
243                ssh_log_out = sshlib.read_command_output()
244                BuiltIn().log_to_console(ssh_log_out)
245            else:
246                if open_connection_args['alias'] == "device_connection":
247                    stdout = sshlib.write(cmd_buf)
248                    stderr = ""
249                    rc = 0
250                else:
251                    stdout, stderr, rc = \
252                        func_timer.run(sshlib.execute_command,
253                                       cmd_buf,
254                                       return_stdout=True,
255                                       return_stderr=True,
256                                       return_rc=True,
257                                       time_out=time_out)
258                    BuiltIn().log_to_console(stdout)
259        except Exception:
260            except_type, except_value, except_traceback = sys.exc_info()
261            gp.lprint_var(except_type)
262            gp.lprint_varx("except_value", str(except_value))
263            # This may be our last time through the retry loop, so setting
264            # return variables.
265            rc = 1
266            stderr = str(except_value)
267            stdout = ""
268
269            if except_type is exceptions.AssertionError and\
270               re.match(r"Connection not open", str(except_value)):
271                try:
272                    login_ssh(login_args)
273                    # Now we must continue to next loop iteration to retry the
274                    # execute_command.
275                    continue
276                except Exception:
277                    except_type, except_value, except_traceback =\
278                        sys.exc_info()
279                    rc = 1
280                    stderr = str(except_value)
281                    stdout = ""
282                    break
283
284            if (except_type is paramiko.ssh_exception.SSHException
285                and re.match(r"SSH session not active", str(except_value))) or\
286               ((except_type is socket.error
287                 or except_type is ConnectionResetError)
288                and re.match(r"\[Errno 104\] Connection reset by peer",
289                             str(except_value))) or\
290               (except_type is paramiko.ssh_exception.SSHException
291                and re.match(r"Timeout opening channel\.",
292                             str(except_value))):
293                # Close and re-open a connection.
294                # Note: close_connection() doesn't appear to get rid of the
295                # connection.  It merely closes it.  Since there is a concern
296                # about over-consumption of resources, we use
297                # close_all_connections() which also gets rid of all
298                # connections.
299                gp.lprint_timen("Closing all connections.")
300                sshlib.close_all_connections()
301                gp.lprint_timen("Connecting to "
302                                + open_connection_args['host'] + ".")
303                cix = sshlib.open_connection(**open_connection_args)
304                login_ssh(login_args)
305                continue
306
307            # We do not handle any other RuntimeErrors so we will raise the exception again.
308            sshlib.close_all_connections()
309            gp.lprintn(traceback.format_exc())
310            raise(except_value)
311
312        # If we get to this point, the command was executed.
313        break
314
315    if fork:
316        return
317
318    if rc != 0 and print_err:
319        gp.print_var(rc, gp.hexa())
320        if not print_out:
321            gp.print_var(stderr)
322            gp.print_var(stdout)
323
324    if print_out:
325        gp.printn(stderr + stdout)
326
327    if not ignore_err:
328        message = gp.sprint_error("The prior SSH"
329                                  + " command returned a non-zero return"
330                                  + " code:\n"
331                                  + gp.sprint_var(rc, gp.hexa()) + stderr
332                                  + "\n")
333        BuiltIn().should_be_equal(rc, 0, message)
334
335    if open_connection_args['alias'] == "device_connection":
336        return stdout
337    return stdout, stderr, rc
338