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