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