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