system/TestSsh.py

   1 # Thierry Parmentelat <thierry.parmentelat@inria.fr>
   2 # Copyright (C) 2010 INRIA
   3 #
   4 # class for issuing commands on a box, either local or remote
   5 #
   6 # the notion of 'buildname' is for providing each test run with a dir of its own
   7 # buildname is generally the name of the build being tested, and can be considered unique
   8 #
   9 # thus 'run_in_buildname' mostly :
  10 # (*) either runs locally in . - as on a local node we are already in a dedicated directory
  11 # (*) or makes sure that there's a remote dir called 'buildname' and runs in it
  12 #
  13 # also, the copy operations
  14 # (*) either do nothing if ran locally
  15 # (*) or copy a local file into the remote 'buildname'
  16 #
  17
  18 import sys
  19 import os.path
  20 import utils
  21 import shutil
  22
  23 class TestSsh:
  24
  25     # inserts a backslash before each occurence of the following chars
  26     # \ " ' < > & | ; ( ) $ * ~
  27     @staticmethod
  28     def backslash_shell_specials(command):
  29         result = ''
  30         for char in command:
  31             if char in "\\\"'<>&|;()$*~":
  32                 result += '\\'+char
  33             else:
  34                 result += char
  35         return result
  36
  37     # check main IP address against the provided hostname
  38     @staticmethod
  39     def is_local_hostname(hostname):
  40         if hostname == "localhost":
  41             return True
  42         import socket
  43         try:
  44             local_ip = socket.gethostbyname(socket.gethostname())
  45             remote_ip = socket.gethostbyname(hostname)
  46             return local_ip == remote_ip
  47         except:
  48             utils.header("WARNING : something wrong in is_local_hostname with hostname={}".format(hostname))
  49             return False
  50
  51     # some boxes have their working space in user's homedir (/root),
  52     # some others in a dedicated area with max. space (/vservers)
  53     # when root is not specified we use the homedir
  54     def __init__(self, hostname, buildname=None, key=None, username=None, unknown_host=True, root=None):
  55         self.hostname = hostname
  56         self.buildname = buildname
  57         self.key = key
  58         self.username = username
  59         self.unknown_host = unknown_host
  60         self.root = root
  61
  62     def __repr__(self):
  63         result = "{}@{}".format(self.username or 'root', self.hostname)
  64         if self.key:
  65             result += " <key {}>".format(self.key)
  66         return result
  67
  68     def is_local(self):
  69         return TestSsh.is_local_hostname(self.hostname)
  70
  71     std_options="-o BatchMode=yes -o StrictHostKeyChecking=no -o CheckHostIP=no -o ConnectTimeout=5 "
  72     unknown_option="-o UserKnownHostsFile=/dev/null "
  73
  74     def key_part(self):
  75         if not self.key:
  76             return ""
  77         return "-i {} ".format(self.key)
  78
  79     def hostname_part(self):
  80         if not self.username:
  81             return self.hostname
  82         else:
  83             return "{}@{}".format(self.username,self.hostname)
  84
  85     # command gets run on the right box
  86     def actual_command(self, command, keep_stdin=False, dry_run=False, backslash=True):
  87         if self.is_local():
  88             return command
  89         ssh_command = "ssh "
  90         if not dry_run:
  91             if not keep_stdin:
  92                 ssh_command += "-n "
  93             ssh_command += TestSsh.std_options
  94             if self.unknown_host: ssh_command += TestSsh.unknown_option
  95         ssh_command += self.key_part()
  96         ssh_command += self.hostname_part() + " "
  97         if backslash:
  98             ssh_command += TestSsh.backslash_shell_specials(command)
  99         else:
 100             ssh_command += command
 101         return ssh_command
 102
 103     # same in argv form
 104     def actual_argv (self, argv, keep_stdin=False, dry_run=False):
 105         if self.is_local():
 106             return argv
 107         ssh_argv = []
 108         ssh_argv.append('ssh')
 109         if not dry_run:
 110             if not keep_stdin:
 111                 ssh_argv.append('-n')
 112             ssh_argv += TestSsh.std_options.split()
 113             if self.unknown_host:
 114                 ssh_argv += TestSsh.unknown_option.split()
 115         ssh_argv += self.key_part().split()
 116         ssh_argv.append(self.hostname_part())
 117         ssh_argv += argv
 118         return ssh_argv
 119
 120     def header(self, message):
 121         if not message:
 122             return
 123         print("===============",message)
 124         sys.stdout.flush()
 125
 126     def run(self, command, message=None, background=False, dry_run=False):
 127         local_command = self.actual_command(command, dry_run=dry_run)
 128         if dry_run:
 129             utils.header("DRY RUN " + local_command)
 130             return 0
 131         else:
 132             self.header(message)
 133             return utils.system(local_command, background)
 134
 135     def run_in_buildname(self, command, background=False, dry_run=False):
 136         if self.is_local():
 137             return utils.system(command, background)
 138         self.create_buildname_once(dry_run)
 139         return self.run("cd {} ; {}".format(self.fullname(self.buildname), command),
 140                         background=background, dry_run=dry_run)
 141
 142     def fullname(self, dirname):
 143         if self.root==None:
 144             return dirname
 145         else:
 146             return os.path.join(self.root,dirname)
 147
 148     def mkdir(self, dirname=None, abs=False, dry_run=False):
 149         if self.is_local():
 150             if dirname:
 151                 return os.path.mkdir(dirname)
 152             return 0
 153         # ab. paths remain as-is
 154         if not abs:
 155             if dirname:
 156                 dirname = "{}/{}".format(self.buildname, dirname)
 157             else:
 158                 dirname = self.buildname
 159             dirname = self.fullname(dirname)
 160         if dirname == '.':
 161             return
 162         return self.run("mkdir -p {}".format(dirname), dry_run=dry_run)
 163
 164     def rmdir(self, dirname=None, dry_run=False):
 165         if self.is_local():
 166             if dirname:
 167                 return shutil.rmtree(dirname)
 168             return 0
 169         if dirname:
 170             dirname = "{}/{}".format(self.buildname, dirname)
 171         else:
 172             dirname = self.buildname
 173         dirname = self.fullname(dirname)
 174         return self.run("rm -rf {}".format(dirname), dry_run=dry_run)
 175
 176     def create_buildname_once(self, dry_run):
 177         if self.is_local():
 178             return
 179         # create remote buildname on demand
 180         try:
 181             self.buildname_created
 182         except:
 183             self.mkdir(dry_run=dry_run)
 184             self.buildname_created = True
 185
 186     def copy(self, local_file, recursive=False, dry_run=False):
 187         if self.is_local():
 188             return 0
 189         self.create_buildname_once(dry_run)
 190         scp_command = "scp "
 191         if not dry_run:
 192             scp_command += TestSsh.std_options
 193         if recursive:
 194             scp_command += "-r "
 195         scp_command += self.key_part()
 196         scp_command += "{} {}:{}/{}".format(local_file, self.hostname_part(),
 197                                             self.fullname(self.buildname),
 198                                             os.path.basename(local_file) or ".")
 199         if dry_run:
 200             utils.header("DRY RUN TestSsh.copy {}".format(scp_command))
 201             # need to be consistent with the non-dry-run mode
 202             return 0
 203         return utils.system(scp_command)
 204
 205     def copy_abs(self, local_file, remote_file,
 206                  recursive=False, dry_run=False):
 207         if self.is_local():
 208             dest = ""
 209         else:
 210             dest = "{}:".format(self.hostname_part())
 211         scp_command = "scp "
 212         scp_command += TestSsh.std_options
 213         if recursive:
 214             scp_command += "-r "
 215         scp_command += self.key_part()
 216         scp_command += "{} {}{}".format(local_file, dest, remote_file)
 217         if dry_run:
 218             utils.header("DRY RUN TestSsh.copy {}".format(scp_command))
 219             # need to be consistent with the non-dry-run mode
 220             return 0
 221         return utils.system(scp_command)
 222
 223     def copy_home(self, local_file, recursive=False):
 224         return self.copy_abs(local_file, os.path.basename(local_file), recursive)
 225
 226     def fetch (self, remote_file, local_file, recursive=False, dry_run=False):
 227         if self.is_local():
 228             command="cp "
 229             if recursive:
 230                 command += "-r "
 231             command += "{} {}".format(remote_file, local_file)
 232         else:
 233             command = "scp "
 234             if not dry_run:
 235                 command += TestSsh.std_options
 236             if recursive:
 237                 command += "-r "
 238             command += self.key_part()
 239             # absolute path - do not preprend buildname
 240             if remote_file.find("/") == 0:
 241                 remote_path = remote_file
 242             else:
 243                 remote_path = "{}/{}".format(self.buildname, remote_file)
 244                 remote_path = self.fullname(remote_path)
 245             command += "{}:{} {}".format(self.hostname_part(), remote_path, local_file)
 246         return utils.system(command)
 247
 248     # this is only to avoid harmless message when host cannot be identified
 249     # convenience only
 250     # the only place where this is needed is when tring to reach a slice in a node,
 251     # which is done from the test master box
 252     def clear_known_hosts(self):
 253         known_hosts = "{}/.ssh/known_hosts".format(os.getenv("HOME"))
 254         utils.header("Clearing entry for {} in {}".format(self.hostname, known_hosts))
 255         return utils.system("sed -i -e /^{}/d {}".format(self.hostname, known_hosts))
 256