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