vstf/vstf/common/ssh.py

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