vstf/vstf/common/ssh.py

   1 '''
   2 Created on 2015-7-23
   3
   4 @author: y00228926
   5 '''
   6 import os
   7 import logging
   8 from stat import S_ISDIR
   9 import Queue
  10 import shutil
  11 import paramiko
  12 from paramiko.ssh_exception import AuthenticationException
  13
  14 LOG = logging.getLogger(__name__)
  15
  16
  17 class SSHClientContext(paramiko.SSHClient):
  18     def __init__(self, ip, user, passwd, port=22):
  19         self.host = ip
  20         self.user = user
  21         self.passwd = passwd
  22         self.port = port
  23         super(SSHClientContext, self).__init__()
  24
  25     def sync_exec_command(self, cmd):
  26         _, stdout, stderr = self.exec_command(cmd)
  27         ret = stdout.channel.recv_exit_status()
  28         out = stdout.read().strip()
  29         err = stderr.read().strip()
  30         LOG.info("in %s,%s,return:%s,output:%s:error:%s" % (self.host, cmd, ret, out, err))
  31         return ret, out, err
  32
  33     def connect(self):
  34         super(SSHClientContext, self).connect(self.host, self.port, self.user, self.passwd, timeout=10)
  35
  36     def __enter__(self):
  37         self.set_missing_host_key_policy(paramiko.AutoAddPolicy())
  38         return self
  39
  40     def __exit__(self, exc_type, exc_val, exc_tb):
  41         self.close()
  42         if exc_type == AuthenticationException:
  43             return False
  44
  45
  46 class SFTPClientContext(object):
  47     def __init__(self, ip, user, passwd, port=22):
  48         self.host = ip
  49         self.passwd = passwd
  50         self.user = user
  51         self.port = port
  52
  53     def connect(self):
  54         self.t = paramiko.Transport((self.host, self.port))
  55         self.t.connect(username=self.user, password=self.passwd)
  56         self.sftp = paramiko.SFTPClient.from_transport(self.t)
  57
  58     def get(self, remote, local):
  59         self.sftp.get(remote, local)
  60
  61     def put(self, local, remote):
  62         self.sftp.put(local, remote)
  63
  64     def mkdir(self, path):
  65         self.sftp.mkdir(path)
  66
  67     def rmdir(self, path):
  68         self.sftp.rmdir(path)
  69
  70     def __enter__(self):
  71         return self
  72
  73     def __exit__(self, exc_type, exc_val, exc_tb):
  74         if exc_type == TypeError:
  75             return False
  76         return False
  77
  78
  79 def upload_conf_file(host, user, passwd, src, dst):
  80     with SFTPClientContext(host, user, passwd) as ftp:
  81         ftp.connect()
  82         LOG.info('putting file:%s to %s:%s' % (src, host, dst))
  83         ftp.put(src, dst)
  84
  85
  86 def upload_dir(host, user, passwd, local_dir, remote_dir):
  87     assert remote_dir.startswith('/')
  88     assert local_dir != '/'
  89     while local_dir.endswith('/'):
  90         local_dir = local_dir[:-1]
  91     while remote_dir.endswith('/'):
  92         remote_dir = remote_dir[:-1]
  93     remote_dir = os.path.join(remote_dir, os.path.basename(local_dir))
  94     ret, _, _ = run_cmd(host, user, passwd, "sudo rm -rf %s" % remote_dir)
  95     if ret != 0 and ret != 1:
  96         LOG.error("somehow failed in rm -rf %s on host:%s,return:%s" % (remote_dir, host, ret))
  97         exit(1)
  98     with SFTPClientContext(host, user, passwd) as sftp:
  99         sftp.connect()
 100         for root, dirs, files in os.walk(local_dir):
 101             for filename in files:
 102                 local_file = os.path.join(root, filename)
 103                 remote_file = local_file.replace(local_dir, remote_dir)
 104                 try:
 105                     sftp.put(local_file, remote_file)
 106                 except IOError:
 107                     sftp.mkdir(os.path.split(remote_file)[0])
 108                     sftp.put(local_file, remote_file)
 109                 LOG.info("upload %s to remote %s" % (local_file, remote_file))
 110             for name in dirs:
 111                 local_path = os.path.join(root, name)
 112                 remote_path = local_path.replace(local_dir, remote_dir)
 113                 try:
 114                     sftp.mkdir(remote_path)
 115                     LOG.info("mkdir path %s" % remote_path)
 116                 except Exception, e:
 117                     raise
 118     return remote_dir
 119
 120
 121 def isdir(path, sftp):
 122     exists = True
 123     is_dir = False
 124     file_stat = None
 125     try:
 126         file_stat = sftp.stat(path).st_mode
 127         is_dir = S_ISDIR(file_stat)
 128     except IOError:
 129         exists = False
 130     return exists, is_dir, file_stat
 131
 132
 133 def download_file(host, user, passwd, remote_path, local_path):
 134     assert not remote_path.endswith('/')
 135     remote_file_name = os.path.basename(remote_path)
 136     if local_path.endswith('/'):
 137         if not os.path.exists(local_path):
 138             raise Exception('path:%s not exist.' % local_path)
 139         dest = os.path.join(local_path, remote_file_name)
 140     else:
 141         if os.path.isdir(local_path):
 142             dest = os.path.join(local_path, remote_file_name)
 143         else:
 144             dir_path = os.path.dirname(local_path)
 145             if not os.path.exists(dir_path):
 146                 raise Exception('path:%s not exist' % dir_path)
 147             dest = local_path
 148     transport = paramiko.Transport((host, 22))
 149     transport.connect(username=user, password=passwd)
 150     sftp = paramiko.SFTPClient.from_transport(transport)
 151     exists, is_dir, st = isdir(remote_path, sftp)
 152     if exists and not is_dir:
 153         sftp.get(remote_path, dest)
 154         os.chmod(dest, st)
 155     else:
 156         raise Exception('error:cannot find the file or file is dir')
 157     return True
 158
 159
 160 def download_dir(host, user, passwd, remote_path, local_path):
 161     while remote_path.endswith('/'):
 162         remote_path = remote_path[:-1]
 163     if local_path.endswith('/'):
 164         if not os.path.exists(local_path):
 165             raise Exception('path:%s not exist.' % local_path)
 166         dest_path = os.path.join(local_path, os.path.basename(remote_path))
 167     else:
 168         if os.path.isdir(local_path):
 169             dest_path = os.path.join(local_path, os.path.basename(remote_path))
 170         else:
 171             dir_name = os.path.dirname(local_path)
 172             if os.path.exists(dir_name):
 173                 dest_path = local_path
 174             else:
 175                 raise Exception('path:%s is not exists' % dir_name)
 176     LOG.info("download_dir from host:%s:%s to dest:%s" % (host, remote_path, dest_path))
 177     transport = paramiko.Transport((host, 22))
 178     transport.connect(username=user, password=passwd)
 179     sftp = paramiko.SFTPClient.from_transport(transport)
 180     exists, is_dir, _ = isdir(remote_path, sftp)
 181     if exists and is_dir:
 182         q = Queue.Queue(0)
 183         q.put(remote_path)
 184         while not q.empty():
 185             path = q.get()
 186             st = sftp.lstat(path).st_mode
 187             relative_path = path[len(remote_path):]
 188             if relative_path.startswith('/'): relative_path = relative_path[1:]
 189             local = os.path.join(dest_path, relative_path)
 190             if os.path.exists(local):
 191                 shutil.rmtree(local)
 192             os.mkdir(local)
 193             os.chmod(local, st)
 194             file_list = sftp.listdir(path)
 195             for item in file_list:
 196                 fullpath = os.path.join(path, item)
 197                 _, is_dir, st = isdir(fullpath, sftp)
 198                 if is_dir:
 199                     q.put(fullpath)
 200                 else:
 201                     dest = os.path.join(local, item)
 202                     sftp.get(fullpath, dest)
 203                     os.chmod(dest, st)
 204     else:
 205         raise Exception('path:%s:%s not exists or is not a dir' % (host, remote_path))
 206     return dest_path
 207
 208
 209 def run_cmd(host, user, passwd, cmd):
 210     with SSHClientContext(host, user, passwd) as ssh:
 211         ssh.connect()
 212         ret, stdout, stderr = ssh.sync_exec_command(cmd)
 213     return ret, stdout, stderr
 214
 215
 216 class SshFileTransfer(object):
 217     def __init__(self, ip, user, passwd):
 218         self.ip, self.user, self.passwd = ip, user, passwd
 219
 220     def upload_dir(self, src, dst):
 221         return upload_dir(self.ip, self.user, self.passwd, src, dst)
 222
 223     def download_dir(self, src, dst):
 224         download_dir(self.ip, self.user, self.passwd, src, dst)
 225
 226     def upload_file(self, src, dst):
 227         upload_conf_file(self.ip, self.user, self.passwd, src, dst)
 228
 229     def download_file(self, src, dst):
 230         download_file(self.ip, self.user, self.passwd, src, dst)