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