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