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