src/tools/telemetry/telemetry/core/backends/chrome/cros_interface.py

   1 # Copyright 2013 The Chromium Authors. All rights reserved.
   2 # Use of this source code is governed by a BSD-style license that can be
   3 # found in the LICENSE file.
   4 """A wrapper around ssh for common operations on a CrOS-based device"""
   5 import logging
   6 import os
   7 import re
   8 import shutil
   9 import subprocess
  10 import sys
  11 import tempfile
  12
  13 # TODO(nduca): This whole file is built up around making individual ssh calls
  14 # for each operation. It really could get away with a single ssh session built
  15 # around pexpect, I suspect, if we wanted it to be faster. But, this was
  16 # convenient.
  17
  18 def IsRunningOnCrosDevice():
  19   """Returns True if we're on a ChromeOS device."""
  20   lsb_release = '/etc/lsb-release'
  21   if sys.platform.startswith('linux') and os.path.exists(lsb_release):
  22     with open(lsb_release, 'r') as f:
  23       res = f.read()
  24       if res.count('CHROMEOS_RELEASE_NAME'):
  25         return True
  26   return False
  27
  28 def RunCmd(args, cwd=None, quiet=False):
  29   """Opens a subprocess to execute a program and returns its return value.
  30
  31   Args:
  32     args: A string or a sequence of program arguments. The program to execute is
  33       the string or the first item in the args sequence.
  34     cwd: If not None, the subprocess's current directory will be changed to
  35       |cwd| before it's executed.
  36
  37   Returns:
  38     Return code from the command execution.
  39   """
  40   if not quiet:
  41     logging.debug(' '.join(args) + ' ' + (cwd or ''))
  42   with open(os.devnull, 'w') as devnull:
  43     p = subprocess.Popen(args=args, cwd=cwd, stdout=devnull,
  44                          stderr=devnull, stdin=devnull, shell=False)
  45     return p.wait()
  46
  47 def GetAllCmdOutput(args, cwd=None, quiet=False):
  48   """Open a subprocess to execute a program and returns its output.
  49
  50   Args:
  51     args: A string or a sequence of program arguments. The program to execute is
  52       the string or the first item in the args sequence.
  53     cwd: If not None, the subprocess's current directory will be changed to
  54       |cwd| before it's executed.
  55
  56   Returns:
  57     Captures and returns the command's stdout.
  58     Prints the command's stderr to logger (which defaults to stdout).
  59   """
  60   if not quiet:
  61     logging.debug(' '.join(args) + ' ' + (cwd or ''))
  62   with open(os.devnull, 'w') as devnull:
  63     p = subprocess.Popen(args=args, cwd=cwd, stdout=subprocess.PIPE,
  64                          stderr=subprocess.PIPE, stdin=devnull)
  65     stdout, stderr = p.communicate()
  66     if not quiet:
  67       logging.debug(' > stdout=[%s], stderr=[%s]', stdout, stderr)
  68     return stdout, stderr
  69
  70 def HasSSH():
  71   try:
  72     RunCmd(['ssh'], quiet=True)
  73     RunCmd(['scp'], quiet=True)
  74     logging.debug("HasSSH()->True")
  75     return True
  76   except OSError:
  77     logging.debug("HasSSH()->False")
  78     return False
  79
  80 class LoginException(Exception):
  81   pass
  82
  83 class KeylessLoginRequiredException(LoginException):
  84   pass
  85
  86 class CrOSInterface(object):
  87   # pylint: disable=R0923
  88   def __init__(self, hostname = None, ssh_identity = None):
  89     self._hostname = hostname
  90     # List of ports generated from GetRemotePort() that may not be in use yet.
  91     self._reserved_ports = []
  92
  93     if self.local:
  94       return
  95
  96     self._ssh_identity = None
  97     self._ssh_args = ['-o ConnectTimeout=5',
  98                       '-o StrictHostKeyChecking=no',
  99                       '-o KbdInteractiveAuthentication=no',
 100                       '-o PreferredAuthentications=publickey',
 101                       '-o UserKnownHostsFile=/dev/null']
 102
 103     if ssh_identity:
 104       self._ssh_identity = os.path.abspath(os.path.expanduser(ssh_identity))
 105
 106   @property
 107   def local(self):
 108     return not self._hostname
 109
 110   @property
 111   def hostname(self):
 112     return self._hostname
 113
 114   def FormSSHCommandLine(self, args, extra_ssh_args=None):
 115     if self.local:
 116       # We run the command through the shell locally for consistency with
 117       # how commands are run through SSH (crbug.com/239161). This work
 118       # around will be unnecessary once we implement a persistent SSH
 119       # connection to run remote commands (crbug.com/239607).
 120       return ['sh', '-c', " ".join(args)]
 121
 122     full_args = ['ssh',
 123                  '-o ForwardX11=no',
 124                  '-o ForwardX11Trusted=no',
 125                  '-n'] + self._ssh_args
 126     if self._ssh_identity is not None:
 127       full_args.extend(['-i', self._ssh_identity])
 128     if extra_ssh_args:
 129       full_args.extend(extra_ssh_args)
 130     full_args.append('root@%s' % self._hostname)
 131     full_args.extend(args)
 132     return full_args
 133
 134   def _RemoveSSHWarnings(self, toClean):
 135     """Removes specific ssh warning lines from a string.
 136
 137     Args:
 138       toClean: A string that may be containing multiple lines.
 139
 140     Returns:
 141       A copy of toClean with all the Warning lines removed.
 142     """
 143     # Remove the Warning about connecting to a new host for the first time.
 144     return re.sub('Warning: Permanently added [^\n]* to the list of known '
 145                   'hosts.\s\n', '', toClean)
 146
 147   def RunCmdOnDevice(self, args, cwd=None, quiet=False):
 148     stdout, stderr = GetAllCmdOutput(
 149         self.FormSSHCommandLine(args), cwd, quiet=quiet)
 150     # The initial login will add the host to the hosts file but will also print
 151     # a warning to stderr that we need to remove.
 152     stderr = self._RemoveSSHWarnings(stderr)
 153     return stdout, stderr
 154
 155   def TryLogin(self):
 156     logging.debug('TryLogin()')
 157     assert not self.local
 158     stdout, stderr = self.RunCmdOnDevice(['echo', '$USER'], quiet=True)
 159     if stderr != '':
 160       if 'Host key verification failed' in stderr:
 161         raise LoginException(('%s host key verification failed. ' +
 162                              'SSH to it manually to fix connectivity.') %
 163             self._hostname)
 164       if 'Operation timed out' in stderr:
 165         raise LoginException('Timed out while logging into %s' % self._hostname)
 166       if 'UNPROTECTED PRIVATE KEY FILE!' in stderr:
 167         raise LoginException('Permissions for %s are too open. To fix this,\n'
 168                              'chmod 600 %s' % (self._ssh_identity,
 169                                                self._ssh_identity))
 170       if 'Permission denied (publickey,keyboard-interactive)' in stderr:
 171         raise KeylessLoginRequiredException(
 172           'Need to set up ssh auth for %s' % self._hostname)
 173       raise LoginException('While logging into %s, got %s' % (
 174           self._hostname, stderr))
 175     if stdout != 'root\n':
 176       raise LoginException(
 177         'Logged into %s, expected $USER=root, but got %s.' % (
 178           self._hostname, stdout))
 179
 180   def FileExistsOnDevice(self, file_name):
 181     if self.local:
 182       return os.path.exists(file_name)
 183
 184     stdout, stderr = self.RunCmdOnDevice([
 185         'if', 'test', '-e', file_name, ';',
 186         'then', 'echo', '1', ';',
 187         'fi'
 188         ], quiet=True)
 189     if stderr != '':
 190       if "Connection timed out" in stderr:
 191         raise OSError('Machine wasn\'t responding to ssh: %s' %
 192                       stderr)
 193       raise OSError('Unepected error: %s' % stderr)
 194     exists = stdout == '1\n'
 195     logging.debug("FileExistsOnDevice(<text>, %s)->%s" % (file_name, exists))
 196     return exists
 197
 198   def PushFile(self, filename, remote_filename):
 199     if self.local:
 200       args = ['cp', '-r', filename, remote_filename]
 201       stdout, stderr = GetAllCmdOutput(args, quiet=True)
 202       if stderr != '':
 203         raise OSError('No such file or directory %s' % stderr)
 204       return
 205
 206     args = ['scp', '-r' ] + self._ssh_args
 207     if self._ssh_identity:
 208       args.extend(['-i', self._ssh_identity])
 209
 210     args.extend([os.path.abspath(filename),
 211                  'root@%s:%s' % (self._hostname, remote_filename)])
 212
 213     stdout, stderr = GetAllCmdOutput(args, quiet=True)
 214     stderr = self._RemoveSSHWarnings(stderr)
 215     if stderr != '':
 216       raise OSError('No such file or directory %s' % stderr)
 217
 218   def PushContents(self, text, remote_filename):
 219     logging.debug("PushContents(<text>, %s)" % remote_filename)
 220     with tempfile.NamedTemporaryFile() as f:
 221       f.write(text)
 222       f.flush()
 223       self.PushFile(f.name, remote_filename)
 224
 225   def GetFile(self, filename, destfile=None):
 226     """Copies a local file |filename| to |destfile| on the device.
 227
 228     Args:
 229       filename: The name of the local source file.
 230       destfile: The name of the file to copy to, and if it is not specified
 231         then it is the basename of the source file.
 232
 233     """
 234     logging.debug("GetFile(%s, %s)" % (filename, destfile))
 235     if self.local:
 236       if destfile is not None and destfile != filename:
 237         shutil.copyfile(filename, destfile)
 238       return
 239
 240     if destfile is None:
 241       destfile = os.path.basename(filename)
 242     args = ['scp'] + self._ssh_args
 243     if self._ssh_identity:
 244       args.extend(['-i', self._ssh_identity])
 245
 246     args.extend(['root@%s:%s' % (self._hostname, filename),
 247                  os.path.abspath(destfile)])
 248     stdout, stderr = GetAllCmdOutput(args, quiet=True)
 249     stderr = self._RemoveSSHWarnings(stderr)
 250     if stderr != '':
 251       raise OSError('No such file or directory %s' % stderr)
 252
 253   def GetFileContents(self, filename):
 254     """Get the contents of a file on the device.
 255
 256     Args:
 257       filename: The name of the file on the device.
 258
 259     Returns:
 260       A string containing the contents of the file.
 261     """
 262     # TODO: handle the self.local case
 263     assert not self.local
 264     t = tempfile.NamedTemporaryFile()
 265     self.GetFile(filename, t.name)
 266     with open(t.name, 'r') as f2:
 267       res = f2.read()
 268       logging.debug("GetFileContents(%s)->%s" % (filename, res))
 269       f2.close()
 270       return res
 271
 272   def ListProcesses(self):
 273     """Returns (pid, cmd, ppid, state) of all processes on the device."""
 274     stdout, stderr = self.RunCmdOnDevice([
 275         '/bin/ps', '--no-headers',
 276         '-A',
 277         '-o', 'pid,ppid,args:4096,state'], quiet=True)
 278     assert stderr == '', stderr
 279     procs = []
 280     for l in stdout.split('\n'): # pylint: disable=E1103
 281       if l == '':
 282         continue
 283       m = re.match('^\s*(\d+)\s+(\d+)\s+(.+)\s+(.+)', l, re.DOTALL)
 284       assert m
 285       procs.append((int(m.group(1)), m.group(3).rstrip(),
 286                     int(m.group(2)), m.group(4)))
 287     logging.debug("ListProcesses(<predicate>)->[%i processes]" % len(procs))
 288     return procs
 289
 290   def RmRF(self, filename):
 291     logging.debug("rm -rf %s" % filename)
 292     self.RunCmdOnDevice(['rm', '-rf', filename], quiet=True)
 293
 294   def Chown(self, filename):
 295     self.RunCmdOnDevice(['chown', '-R', 'chronos:chronos', filename])
 296
 297   def KillAllMatching(self, predicate):
 298     kills = ['kill', '-KILL']
 299     for pid, cmd, _, _ in self.ListProcesses():
 300       if predicate(cmd):
 301         logging.info('Killing %s, pid %d' % cmd, pid)
 302         kills.append(pid)
 303     logging.debug("KillAllMatching(<predicate>)->%i" % (len(kills) - 2))
 304     if len(kills) > 2:
 305       self.RunCmdOnDevice(kills, quiet=True)
 306     return len(kills) - 2
 307
 308   def IsServiceRunning(self, service_name):
 309     stdout, stderr = self.RunCmdOnDevice([
 310         'status', service_name], quiet=True)
 311     assert stderr == '', stderr
 312     running = 'running, process' in stdout
 313     logging.debug("IsServiceRunning(%s)->%s" % (service_name, running))
 314     return running
 315
 316   def GetRemotePort(self):
 317     netstat = self.RunCmdOnDevice(['netstat', '-ant'])
 318     netstat = netstat[0].split('\n')
 319     ports_in_use = []
 320
 321     for line in netstat[2:]:
 322       if not line:
 323         continue
 324       address_in_use = line.split()[3]
 325       port_in_use = address_in_use.split(':')[-1]
 326       ports_in_use.append(int(port_in_use))
 327
 328     ports_in_use.extend(self._reserved_ports)
 329
 330     new_port = sorted(ports_in_use)[-1] + 1
 331     self._reserved_ports.append(new_port)
 332
 333     return new_port
 334
 335   def IsHTTPServerRunningOnPort(self, port):
 336     wget_output = self.RunCmdOnDevice(
 337         ['wget', 'localhost:%i' % (port), '-T1', '-t1'])
 338
 339     if 'Connection refused' in wget_output[1]:
 340       return False
 341
 342     return True
 343
 344   def FilesystemMountedAt(self, path):
 345     """Returns the filesystem mounted at |path|"""
 346     df_out, _ = self.RunCmdOnDevice(['/bin/df', path])
 347     df_ary = df_out.split('\n')
 348     # 3 lines for title, mount info, and empty line.
 349     if len(df_ary) == 3:
 350       line_ary = df_ary[1].split()
 351       if line_ary:
 352         return line_ary[0]
 353     return None
 354
 355   def CryptohomePath(self, user):
 356     """Returns the cryptohome mount point for |user|."""
 357     return self.RunCmdOnDevice(
 358         ['cryptohome-path', 'user', "'%s'" % user])[0].strip()
 359
 360   def IsCryptohomeMounted(self, username):
 361     """Returns True iff |user|'s cryptohome is mounted."""
 362     profile_path = self.CryptohomePath(username)
 363     mount = self.FilesystemMountedAt(profile_path)
 364     mount_prefix = 'guestfs' if username == '$guest' else '/home/.shadow/'
 365     return mount and mount.startswith(mount_prefix)
 366
 367   def TakeScreenShot(self, screenshot_prefix):
 368     """Takes a screenshot, useful for debugging failures."""
 369     # TODO(achuith): Find a better location for screenshots. Cros autotests
 370     # upload everything in /var/log so use /var/log/screenshots for now.
 371     SCREENSHOT_DIR = '/var/log/screenshots/'
 372     SCREENSHOT_EXT = '.png'
 373
 374     self.RunCmdOnDevice(['mkdir', '-p', SCREENSHOT_DIR])
 375     for i in xrange(25):
 376       screenshot_file = ('%s%s-%d%s' %
 377                          (SCREENSHOT_DIR, screenshot_prefix, i, SCREENSHOT_EXT))
 378       if not self.FileExistsOnDevice(screenshot_file):
 379         self.RunCmdOnDevice([
 380             'DISPLAY=:0.0 XAUTHORITY=/home/chronos/.Xauthority '
 381             '/usr/local/bin/import',
 382             '-window root',
 383             '-depth 8',
 384             screenshot_file])
 385         return
 386     logging.warning('screenshot directory full.')