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