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