1# Copyright 2014 LinkedIn Corp. 2# 3# This file is free software; you can redistribute it and/or 4# modify it under the terms of the GNU Lesser General Public 5# License as published by the Free Software Foundation; either 6# version 2.1 of the License, or (at your option) any later 7# version. 8# 9# This file is distributed in the hope that it will be 10# useful, but WITHOUT ANY WARRANTY; without even the implied 11# warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR 12# PURPOSE. 13# 14# See the GNU Lesser General Public License for more details. 15# You may obtain a copy of the License at 16# https://www.gnu.org/licenses/lgpl-2.1.html 17 18""" 19 20""" 21from contextlib import contextmanager 22import errno 23import logging 24import os 25import re 26import stat 27 28 29logger = logging.getLogger(__name__) 30 31 32class DeploymentError(Exception): 33 """Represents an exception occurring in the deployment module 34 35 Attributes: 36 msg -- explanation of the error 37 """ 38 39 def __init__(self, msg): 40 self.msg = msg 41 42 def __str__(self): 43 return self.msg 44 45 46class ParamikoError(DeploymentError): 47 """Represents an exception if a command Paramiko tries to execute fails 48 49 Attributes: 50 msg -- explanation of the error 51 errors -- a list of lines representing the output to stderr by paramiko 52 """ 53 54 def __init__(self, msg, errors): 55 self.msg = msg 56 self.errors = errors 57 58 def __str__(self): 59 return "{0}\n{1}".format(self.msg, self.errors) 60 61 62def build_os_environment_string(env): 63 """ Creates a string of the form export key0=value0;export key1=value1;... for use in 64 running commands with the specified environment 65 66 :Parameter variables: a dictionay of environmental variables 67 :Returns string: a string that can be prepended to a command to run the command with 68 the environmental variables set 69 """ 70 return "".join(["export {0}={1}; ".format(key, env[key]) for key in env]) 71 72def exec_with_env(ssh, command, msg='', env={}, **kwargs): 73 """ 74 75 :param ssh: 76 :param command: 77 :param msg: 78 :param env: 79 :param synch: 80 :return: 81 """ 82 bash_profile_command = "source .bash_profile > /dev/null 2> /dev/null;" 83 env_command = build_os_environment_string(env) 84 new_command = bash_profile_command + env_command + command 85 if kwargs.get('sync', True): 86 return better_exec_command(ssh, new_command, msg) 87 else: 88 return ssh.exec_command(new_command) 89 90def better_exec_command(ssh, command, msg): 91 """Uses paramiko to execute a command but handles failure by raising a ParamikoError if the command fails. 92 Note that unlike paramiko.SSHClient.exec_command this is not asynchronous because we wait until the exit status is known 93 94 :Parameter ssh: a paramiko SSH Client 95 :Parameter command: the command to execute 96 :Parameter msg: message to print on failure 97 98 :Returns (paramiko.Channel) 99 the underlying channel so that the caller can extract stdout or send to stdin 100 101 :Raises SSHException: if paramiko would raise an SSHException 102 :Raises ParamikoError: if the command produces output to stderr 103 """ 104 chan = ssh.get_transport().open_session() 105 chan.exec_command(command) 106 exit_status = chan.recv_exit_status() 107 if exit_status != 0: 108 msg_str = chan.recv_stderr(1024) 109 err_msgs = [] 110 while len(msg_str) > 0: 111 err_msgs.append(msg_str) 112 msg_str = chan.recv_stderr(1024) 113 err_msg = ''.join(err_msgs) 114 logger.error(err_msg) 115 raise ParamikoError(msg, err_msg) 116 return chan 117 118def log_output(chan): 119 """ 120 logs the output from a remote command 121 the input should be an open channel in the case of synchronous better_exec_command 122 otherwise this will not log anything and simply return to the caller 123 :param chan: 124 :return: 125 """ 126 if hasattr(chan, "recv"): 127 str = chan.recv(1024) 128 msgs = [] 129 while len(str) > 0: 130 msgs.append(str) 131 str = chan.recv(1024) 132 msg = ''.join(msgs).strip() 133 if len(msg) > 0: 134 logger.info(msg) 135 136 137def copy_dir(ftp, filename, outputdir, prefix, pattern=''): 138 """ 139 Recursively copy a directory flattens the output into a single directory but 140 prefixes the files with the path from the original input directory 141 :param ftp: 142 :param filename: 143 :param outputdir: 144 :param prefix: 145 :param pattern: a regex pattern for files to match (by default matches everything) 146 :return: 147 """ 148 try: 149 mode = ftp.stat(filename).st_mode 150 except IOError, e: 151 if e.errno == errno.ENOENT: 152 logger.error("Log file " + filename + " does not exist") 153 pass 154 else: 155 if mode & stat.S_IFREG: 156 if re.match(pattern, filename) is not None: 157 new_file = os.path.join(outputdir, "{0}-{1}".format(prefix, os.path.basename(filename))) 158 ftp.get(filename, new_file) 159 elif mode & stat.S_IFDIR: 160 for f in ftp.listdir(filename): 161 copy_dir(ftp, os.path.join(filename, f), outputdir, 162 "{0}_{1}".format(prefix, os.path.basename(filename)), pattern) 163 164 165@contextmanager 166def open_remote_file(hostname, filename, mode='r', bufsize=-1, username=None, password=None): 167 """ 168 169 :param hostname: 170 :param filename: 171 :return: 172 """ 173 with get_ssh_client(hostname, username=username, password=password) as ssh: 174 sftp = None 175 f = None 176 try: 177 sftp = ssh.open_sftp() 178 f = sftp.open(filename, mode, bufsize) 179 yield f 180 finally: 181 if f is not None: 182 f.close() 183 if sftp is not None: 184 sftp.close() 185 186 187@contextmanager 188def get_sftp_client(hostname, username=None, password=None): 189 with get_ssh_client(hostname, username=username, password=password) as ssh: 190 sftp = None 191 try: 192 sftp = ssh.open_sftp() 193 yield sftp 194 finally: 195 if sftp is not None: 196 sftp.close() 197 198 199@contextmanager 200def get_ssh_client(hostname, username=None, password=None): 201 try: 202 ssh = sshclient() 203 ssh.load_system_host_keys() 204 ssh.connect(hostname, username=username, password=password) 205 yield ssh 206 finally: 207 if ssh is not None: 208 ssh.close() 209 210@contextmanager 211def get_remote_session(hostname, username=None, password=None): 212 with get_ssh_client(hostname, username=username, password=password) as ssh: 213 try: 214 shell = ssh.invoke_shell() 215 yield shell 216 finally: 217 if shell is not None: 218 shell.close() 219 220@contextmanager 221def get_remote_session_with_environment(hostname, env, username=None, password=None): 222 with get_remote_session(hostname, username=username, password=password) as shell: 223 shell.send(build_os_environment_string(env)) 224 shell.send("\n") 225 yield shell 226 227def sshclient(): 228 try: 229 import paramiko 230 ssh = paramiko.SSHClient() 231 ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy()) 232 return ssh 233 except ImportError: 234 return None 235