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