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