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