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            else:
244                if open_connection_args['alias'] == "device_connection":
245                    stdout = sshlib.write(cmd_buf)
246                    stderr = ""
247                    rc = 0
248                else:
249                    stdout, stderr, rc = \
250                        func_timer.run(sshlib.execute_command,
251                                       cmd_buf,
252                                       return_stdout=True,
253                                       return_stderr=True,
254                                       return_rc=True,
255                                       time_out=time_out)
256                    BuiltIn().log_to_console(stdout)
257        except Exception:
258            except_type, except_value, except_traceback = sys.exc_info()
259            gp.lprint_var(except_type)
260            gp.lprint_varx("except_value", str(except_value))
261            # This may be our last time through the retry loop, so setting
262            # return variables.
263            rc = 1
264            stderr = str(except_value)
265            stdout = ""
266
267            if except_type is exceptions.AssertionError and\
268               re.match(r"Connection not open", str(except_value)):
269                try:
270                    login_ssh(login_args)
271                    # Now we must continue to next loop iteration to retry the
272                    # execute_command.
273                    continue
274                except Exception:
275                    except_type, except_value, except_traceback =\
276                        sys.exc_info()
277                    rc = 1
278                    stderr = str(except_value)
279                    stdout = ""
280                    break
281
282            if (except_type is paramiko.ssh_exception.SSHException
283                and re.match(r"SSH session not active", str(except_value))) or\
284               ((except_type is socket.error
285                 or except_type is ConnectionResetError)
286                and re.match(r"\[Errno 104\] Connection reset by peer",
287                             str(except_value))) or\
288               (except_type is paramiko.ssh_exception.SSHException
289                and re.match(r"Timeout opening channel\.",
290                             str(except_value))):
291                # Close and re-open a connection.
292                # Note: close_connection() doesn't appear to get rid of the
293                # connection.  It merely closes it.  Since there is a concern
294                # about over-consumption of resources, we use
295                # close_all_connections() which also gets rid of all
296                # connections.
297                gp.lprint_timen("Closing all connections.")
298                sshlib.close_all_connections()
299                gp.lprint_timen("Connecting to "
300                                + open_connection_args['host'] + ".")
301                cix = sshlib.open_connection(**open_connection_args)
302                login_ssh(login_args)
303                continue
304
305            # We do not handle any other RuntimeErrors so we will raise the exception again.
306            sshlib.close_all_connections()
307            gp.lprintn(traceback.format_exc())
308            raise (except_value)
309
310        # If we get to this point, the command was executed.
311        break
312
313    if fork:
314        return
315
316    if rc != 0 and print_err:
317        gp.print_var(rc, gp.hexa())
318        if not print_out:
319            gp.print_var(stderr)
320            gp.print_var(stdout)
321
322    if print_out:
323        gp.printn(stderr + stdout)
324
325    if not ignore_err:
326        message = gp.sprint_error("The prior SSH"
327                                  + " command returned a non-zero return"
328                                  + " code:\n"
329                                  + gp.sprint_var(rc, gp.hexa()) + stderr
330                                  + "\n")
331        BuiltIn().should_be_equal(rc, 0, message)
332
333    if open_connection_args['alias'] == "device_connection":
334        return stdout
335    return stdout, stderr, rc
336