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