src/third_party/chromite/lib/osutils.py

   1 # Copyright (c) 2012 The Chromium OS 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
   5 """Common file and os related utilities, including tempdir manipulation."""
   6
   7 import collections
   8 import contextlib
   9 import cStringIO
  10 import ctypes
  11 import ctypes.util
  12 import datetime
  13 import errno
  14 import logging
  15 import operator
  16 import os
  17 import pwd
  18 import re
  19 import shutil
  20 import tempfile
  21
  22 from chromite.lib import cros_build_lib
  23 from chromite.lib import retry_util
  24
  25
  26 # Env vars that tempdir can be gotten from; minimally, this
  27 # needs to match python's tempfile module and match normal
  28 # unix standards.
  29 _TEMPDIR_ENV_VARS = ('TMPDIR', 'TEMP', 'TMP')
  30
  31
  32 def GetNonRootUser():
  33   """Returns a non-root user. Defaults to the current user.
  34
  35   If the current user is root, returns the username of the person who
  36   ran the emerge command. If running using sudo, returns the username
  37   of the person who ran the sudo command. If no non-root user is
  38   found, returns None.
  39   """
  40   uid = os.getuid()
  41   if uid == 0:
  42     user = os.environ.get('PORTAGE_USERNAME', os.environ.get('SUDO_USER'))
  43   else:
  44     user = pwd.getpwuid(os.getuid()).pw_name
  45
  46   if user == 'root':
  47     return None
  48   else:
  49     return user
  50
  51
  52 def ExpandPath(path):
  53   """Returns path after passing through realpath and expanduser."""
  54   return os.path.realpath(os.path.expanduser(path))
  55
  56
  57 def WriteFile(path, content, mode='w', atomic=False, makedirs=False):
  58   """Write the given content to disk.
  59
  60   Args:
  61     path: Pathway to write the content to.
  62     content: Content to write.  May be either an iterable, or a string.
  63     mode: Optional; if binary mode is necessary, pass 'wb'.  If appending is
  64           desired, 'w+', etc.
  65     atomic: If the updating of the file should be done atomically.  Note this
  66             option is incompatible w/ append mode.
  67     makedirs: If True, create missing leading directories in the path.
  68   """
  69   write_path = path
  70   if atomic:
  71     write_path = path + '.tmp'
  72
  73   if makedirs:
  74     SafeMakedirs(os.path.dirname(path))
  75
  76   with open(write_path, mode) as f:
  77     f.writelines(cros_build_lib.iflatten_instance(content))
  78
  79   if not atomic:
  80     return
  81
  82   try:
  83     os.rename(write_path, path)
  84   except EnvironmentError:
  85     SafeUnlink(write_path)
  86     raise
  87
  88
  89 def Touch(path, makedirs=False, mode=None):
  90   """Simulate unix touch. Create if doesn't exist and update its timestamp.
  91
  92   Args:
  93     path: a string, file name of the file to touch (creating if not present).
  94     makedirs: If True, create missing leading directories in the path.
  95     mode: The access permissions to set.  In the style of chmod.  Defaults to
  96           using the umask.
  97   """
  98   if makedirs:
  99     SafeMakedirs(os.path.dirname(path))
 100
 101   # Create the file if nonexistant.
 102   open(path, 'a').close()
 103   if mode is not None:
 104     os.chmod(path, mode)
 105   # Update timestamp to right now.
 106   os.utime(path, None)
 107
 108
 109 def ReadFile(path, mode='r'):
 110   """Read a given file on disk.  Primarily useful for one off small files."""
 111   with open(path, mode) as f:
 112     return f.read()
 113
 114
 115 def SafeUnlink(path, sudo=False):
 116   """Unlink a file from disk, ignoring if it doesn't exist.
 117
 118   Returns:
 119     True if the file existed and was removed, False if it didn't exist.
 120   """
 121   if sudo:
 122     try:
 123       cros_build_lib.SudoRunCommand(
 124           ['rm', '--',  path], print_cmd=False, redirect_stderr=True)
 125       return True
 126     except cros_build_lib.RunCommandError:
 127       if os.path.exists(path):
 128         # Technically racey, but oh well; very hard to actually hit...
 129         raise
 130       return False
 131   try:
 132     os.unlink(path)
 133     return True
 134   except EnvironmentError as e:
 135     if e.errno != errno.ENOENT:
 136       raise
 137   return False
 138
 139
 140 def SafeMakedirs(path, mode=0o775, sudo=False, user='root'):
 141   """Make parent directories if needed.  Ignore if existing.
 142
 143   Args:
 144     path: The path to create.  Intermediate directories will be created as
 145           needed.
 146     mode: The access permissions in the style of chmod.
 147     sudo: If True, create it via sudo, thus root owned.
 148     user: If |sudo| is True, run sudo as |user|.
 149
 150   Returns:
 151     True if the directory had to be created, False if otherwise.
 152
 153   Raises:
 154     EnvironmentError: if the makedir failed and it was non sudo.
 155     RunCommandError: If sudo mode, and the command failed for any reason.
 156   """
 157   if sudo:
 158     if os.path.isdir(path):
 159       return False
 160     cros_build_lib.SudoRunCommand(
 161         ['mkdir', '-p', '--mode', oct(mode), path], user=user, print_cmd=False,
 162         redirect_stderr=True, redirect_stdout=True)
 163     return True
 164
 165   try:
 166     os.makedirs(path, mode)
 167     return True
 168   except EnvironmentError as e:
 169     if e.errno != errno.EEXIST or not os.path.isdir(path):
 170       raise
 171
 172   return False
 173
 174
 175 class MakingDirsAsRoot(Exception):
 176   """Raised when creating directories as root."""
 177
 178
 179 def SafeMakedirsNonRoot(path, mode=0o775, user=None):
 180   """Create directories and make sure they are not owned by root.
 181
 182   See SafeMakedirs for the arguments and returns.
 183   """
 184   if user is None:
 185     user = GetNonRootUser()
 186
 187   if user is None or user == 'root':
 188     raise MakingDirsAsRoot('Refusing to create %s as root!' % path)
 189
 190   created = SafeMakedirs(path, mode=mode, user=user)
 191   # Temporary fix: if the directory already exists and is owned by
 192   # root, chown it. This corrects existing root-owned directories.
 193   if not created:
 194     stat_info = os.stat(path)
 195     if stat_info.st_uid == 0:
 196       cros_build_lib.SudoRunCommand(['chown', user, path],
 197                                     print_cmd=False,
 198                                     redirect_stderr=True,
 199                                     redirect_stdout=True)
 200   return created
 201
 202
 203 def RmDir(path, ignore_missing=False, sudo=False):
 204   """Recursively remove a directory.
 205
 206   Args:
 207     path: Path of directory to remove.
 208     ignore_missing: Do not error when path does not exist.
 209     sudo: Remove directories as root.
 210   """
 211   if sudo:
 212     try:
 213       cros_build_lib.SudoRunCommand(
 214           ['rm', '-r%s' % ('f' if ignore_missing else '',), '--', path],
 215           debug_level=logging.DEBUG,
 216           redirect_stdout=True, redirect_stderr=True)
 217     except cros_build_lib.RunCommandError as e:
 218       if not ignore_missing or os.path.exists(path):
 219         # If we're not ignoring the rm ENOENT equivalent, throw it;
 220         # if the pathway still exists, something failed, thus throw it.
 221         raise
 222   else:
 223     try:
 224       shutil.rmtree(path)
 225     except EnvironmentError as e:
 226       if not ignore_missing or e.errno != errno.ENOENT:
 227         raise
 228
 229
 230 def Which(binary, path=None, mode=os.X_OK):
 231   """Return the absolute path to the specified binary.
 232
 233   Args:
 234     binary: The binary to look for.
 235     path: Search path. Defaults to os.environ['PATH'].
 236     mode: File mode to check on the binary.
 237
 238   Returns:
 239     The full path to |binary| if found (with the right mode). Otherwise, None.
 240   """
 241   if path is None:
 242     path = os.environ.get('PATH', '')
 243   for p in path.split(os.pathsep):
 244     p = os.path.join(p, binary)
 245     if os.path.isfile(p) and os.access(p, mode):
 246       return p
 247   return None
 248
 249
 250 def FindDepotTools():
 251   """Returns the location of depot_tools if it is in $PATH."""
 252   gclient_dir = os.path.dirname(Which('gclient.py', mode=os.F_OK) or '')
 253   gitcl_dir = os.path.dirname(Which('git_cl.py', mode=os.F_OK) or '')
 254   if gclient_dir and gclient_dir == gitcl_dir:
 255     return gclient_dir
 256
 257
 258 def FindMissingBinaries(needed_tools):
 259   """Verifies that the required tools are present on the system.
 260
 261   This is especially important for scripts that are intended to run
 262   outside the chroot.
 263
 264   Args:
 265     needed_tools: an array of string specified binaries to look for.
 266
 267   Returns:
 268     If all tools are found, returns the empty list. Otherwise, returns the
 269     list of missing tools.
 270   """
 271   return [binary for binary in needed_tools if Which(binary) is None]
 272
 273
 274 def DirectoryIterator(base_path):
 275   """Iterates through the files and subdirs of a directory."""
 276   for root, dirs, files in os.walk(base_path):
 277     for e in [d + os.sep for d in dirs] + files:
 278       yield os.path.join(root, e)
 279
 280
 281 def IteratePathParents(start_path):
 282   """Generator that iterates through a directory's parents.
 283
 284   Args:
 285     start_path: The path to start from.
 286
 287   Yields:
 288     The passed-in path, along with its parents.  i.e.,
 289     IteratePathParents('/usr/local') would yield '/usr/local', '/usr/', and '/'.
 290   """
 291   path = os.path.abspath(start_path)
 292   yield path
 293   while path.strip('/'):
 294     path = os.path.dirname(path)
 295     yield path
 296
 297
 298 def FindInPathParents(path_to_find, start_path, test_func=None):
 299   """Look for a relative path, ascending through parent directories.
 300
 301   Ascend through parent directories of current path looking for a relative
 302   path.  I.e., given a directory structure like:
 303   -/
 304    |
 305    --usr
 306      |
 307      --bin
 308      |
 309      --local
 310        |
 311        --google
 312
 313   the call FindInPathParents('bin', '/usr/local') would return '/usr/bin', and
 314   the call FindInPathParents('google', '/usr/local') would return
 315   '/usr/local/google'.
 316
 317   Args:
 318     path_to_find: The relative path to look for.
 319     start_path: The path to start the search from.  If |start_path| is a
 320       directory, it will be included in the directories that are searched.
 321     test_func: The function to use to verify the relative path.  Defaults to
 322       os.path.exists.  The function will be passed one argument - the target
 323       path to test.  A True return value will cause AscendingLookup to return
 324       the target.
 325   """
 326   if test_func is None:
 327     test_func = os.path.exists
 328   for path in IteratePathParents(start_path):
 329     target = os.path.join(path, path_to_find)
 330     if test_func(target):
 331       return target
 332   return None
 333
 334
 335 # pylint: disable=W0212,R0904,W0702
 336 def SetGlobalTempDir(tempdir_value, tempdir_env=None):
 337   """Set the global temp directory to the specified |tempdir_value|
 338
 339   Args:
 340     tempdir_value: The new location for the global temp directory.
 341     tempdir_env: Optional. A list of key/value pairs to set in the
 342       environment. If not provided, set all global tempdir environment
 343       variables to point at |tempdir_value|.
 344
 345   Returns:
 346     Returns (old_tempdir_value, old_tempdir_env).
 347
 348     old_tempdir_value: The old value of the global temp directory.
 349     old_tempdir_env: A list of the key/value pairs that control the tempdir
 350       environment and were set prior to this function. If the environment
 351       variable was not set, it is recorded as None.
 352   """
 353   with tempfile._once_lock:
 354     old_tempdir_value = tempfile._get_default_tempdir()
 355     old_tempdir_env = tuple((x, os.environ.get(x)) for x in _TEMPDIR_ENV_VARS)
 356
 357     # Now update TMPDIR/TEMP/TMP, and poke the python
 358     # internals to ensure all subprocess/raw tempfile
 359     # access goes into this location.
 360     if tempdir_env is None:
 361       os.environ.update((x, tempdir_value) for x in _TEMPDIR_ENV_VARS)
 362     else:
 363       for key, value in tempdir_env:
 364         if value is None:
 365           os.environ.pop(key, None)
 366         else:
 367           os.environ[key] = value
 368
 369     # Finally, adjust python's cached value (we know it's cached by here
 370     # since we invoked _get_default_tempdir from above).  Note this
 371     # is necessary since we want *all* output from that point
 372     # forward to go to this location.
 373     tempfile.tempdir = tempdir_value
 374
 375   return (old_tempdir_value, old_tempdir_env)
 376
 377
 378 def _TempDirSetup(self, prefix='tmp', set_global=False, base_dir=None):
 379   """Generate a tempdir, modifying the object, and env to use it.
 380
 381   Specifically, if set_global is True, then from this invocation forward,
 382   python and all subprocesses will use this location for their tempdir.
 383
 384   The matching _TempDirTearDown restores the env to what it was.
 385   """
 386   # Stash the old tempdir that was used so we can
 387   # switch it back on the way out.
 388   self.tempdir = tempfile.mkdtemp(prefix=prefix, dir=base_dir)
 389   os.chmod(self.tempdir, 0o700)
 390
 391   if set_global:
 392     self._orig_tempdir_value, self._orig_tempdir_env = \
 393         SetGlobalTempDir(self.tempdir)
 394
 395
 396 def _TempDirTearDown(self, force_sudo):
 397   # Note that _TempDirSetup may have failed, resulting in these attributes
 398   # not being set; this is why we use getattr here (and must).
 399   tempdir = getattr(self, 'tempdir', None)
 400   try:
 401     if tempdir is not None:
 402       RmDir(tempdir, ignore_missing=True, sudo=force_sudo)
 403   except EnvironmentError as e:
 404     # Suppress ENOENT since we may be invoked
 405     # in a context where parallel wipes of the tempdir
 406     # may be occuring; primarily during hard shutdowns.
 407     if e.errno != errno.ENOENT:
 408       raise
 409
 410   # Restore environment modification if necessary.
 411   orig_tempdir_value = getattr(self, '_orig_tempdir_value', None)
 412   if orig_tempdir_value is not None:
 413     SetGlobalTempDir(orig_tempdir_value, self._orig_tempdir_env)
 414
 415
 416 class TempDir(object):
 417   """Object that creates a temporary directory.
 418
 419   This object can either be used as a context manager or just as a simple
 420   object. The temporary directory is stored as self.tempdir in the object, and
 421   is returned as a string by a 'with' statement.
 422   """
 423
 424   def __init__(self, **kwargs):
 425     """Constructor. Creates the temporary directory.
 426
 427     Args:
 428       prefix: See tempfile.mkdtemp documentation.
 429       base_dir: The directory to place the temporary directory.
 430       set_global: Set this directory as the global temporary directory.
 431       storage: The object that will have its 'tempdir' attribute set.
 432       sudo_rm: Whether the temporary dir will need root privileges to remove.
 433     """
 434     self.kwargs = kwargs.copy()
 435     self.sudo_rm = kwargs.pop('sudo_rm', False)
 436     self.tempdir = None
 437     _TempDirSetup(self, **kwargs)
 438
 439   def Cleanup(self):
 440     """Clean up the temporary directory."""
 441     if self.tempdir is not None:
 442       try:
 443         _TempDirTearDown(self, self.sudo_rm)
 444       finally:
 445         self.tempdir = None
 446
 447   def __enter__(self):
 448     """Return the temporary directory."""
 449     return self.tempdir
 450
 451   def __exit__(self, exc_type, exc_value, exc_traceback):
 452     try:
 453       self.Cleanup()
 454     except Exception:
 455       if exc_type:
 456         # If an exception from inside the context was already in progress,
 457         # log our cleanup exception, then allow the original to resume.
 458         cros_build_lib.Error('While exiting %s:', self, exc_info=True)
 459
 460         if self.tempdir:
 461           # Log all files in tempdir at the time of the failure.
 462           try:
 463             cros_build_lib.Error('Directory contents were:')
 464             for name in os.listdir(self.tempdir):
 465               cros_build_lib.Error('  %s', name)
 466           except OSError:
 467             cros_build_lib.Error('  Directory did not exist.')
 468
 469           # Log all mounts at the time of the failure, since that's the most
 470           # common cause.
 471           mount_results = cros_build_lib.RunCommand(
 472               ['mount'], redirect_stdout=True, combine_stdout_stderr=True,
 473               error_code_ok=True)
 474           cros_build_lib.Error('Mounts were:')
 475           cros_build_lib.Error('  %s', mount_results.output)
 476
 477       else:
 478         # If there was not an exception from the context, raise ours.
 479         raise
 480
 481   def __del__(self):
 482     self.Cleanup()
 483
 484
 485 # pylint: disable=W0212,R0904,W0702
 486 def TempDirDecorator(func):
 487   """Populates self.tempdir with path to a temporary writeable directory."""
 488   def f(self, *args, **kwargs):
 489     with TempDir() as tempdir:
 490       self.tempdir = tempdir
 491       return func(self, *args, **kwargs)
 492
 493   f.__name__ = func.__name__
 494   f.__doc__ = func.__doc__
 495   f.__module__ = func.__module__
 496   return f
 497
 498
 499 def TempFileDecorator(func):
 500   """Populates self.tempfile with path to a temporary writeable file"""
 501   def f(self, *args, **kwargs):
 502     with tempfile.NamedTemporaryFile(dir=self.tempdir, delete=False) as f:
 503       self.tempfile = f.name
 504     return func(self, *args, **kwargs)
 505
 506   f.__name__ = func.__name__
 507   f.__doc__ = func.__doc__
 508   f.__module__ = func.__module__
 509   return TempDirDecorator(f)
 510
 511
 512 # Flags synced from sys/mount.h.  See mount(2) for details.
 513 MS_RDONLY = 1
 514 MS_NOSUID = 2
 515 MS_NODEV = 4
 516 MS_NOEXEC = 8
 517 MS_SYNCHRONOUS = 16
 518 MS_REMOUNT = 32
 519 MS_MANDLOCK = 64
 520 MS_DIRSYNC = 128
 521 MS_NOATIME = 1024
 522 MS_NODIRATIME = 2048
 523 MS_BIND = 4096
 524 MS_MOVE = 8192
 525 MS_REC = 16384
 526 MS_SILENT = 32768
 527 MS_POSIXACL = 1 << 16
 528 MS_UNBINDABLE = 1 << 17
 529 MS_PRIVATE = 1 << 18
 530 MS_SLAVE = 1 << 19
 531 MS_SHARED = 1 << 20
 532 MS_RELATIME = 1 << 21
 533 MS_KERNMOUNT = 1 << 22
 534 MS_I_VERSION =  1 << 23
 535 MS_STRICTATIME = 1 << 24
 536 MS_ACTIVE = 1 << 30
 537 MS_NOUSER = 1 << 31
 538
 539
 540 def Mount(source, target, fstype, flags, data=""):
 541   """Call the mount(2) func; see the man page for details."""
 542   libc = ctypes.CDLL(ctypes.util.find_library('c'), use_errno=True)
 543   if libc.mount(source, target, fstype, ctypes.c_int(flags), data) != 0:
 544     e = ctypes.get_errno()
 545     raise OSError(e, os.strerror(e))
 546
 547
 548 def MountDir(src_path, dst_path, fs_type=None, sudo=True, makedirs=True,
 549              mount_opts=('nodev', 'noexec', 'nosuid'), skip_mtab=False,
 550              **kwargs):
 551   """Mount |src_path| at |dst_path|
 552
 553   Args:
 554     src_path: Source of the new mount.
 555     dst_path: Where to mount things.
 556     fs_type: Specify the filesystem type to use.  Defaults to autodetect.
 557     sudo: Run through sudo.
 558     makedirs: Create |dst_path| if it doesn't exist.
 559     mount_opts: List of options to pass to `mount`.
 560     skip_mtab: Whether to write new entries to /etc/mtab.
 561     kwargs: Pass all other args to RunCommand.
 562   """
 563   if sudo:
 564     runcmd = cros_build_lib.SudoRunCommand
 565   else:
 566     runcmd = cros_build_lib.RunCommand
 567
 568   if makedirs:
 569     SafeMakedirs(dst_path, sudo=sudo)
 570
 571   cmd = ['mount', src_path, dst_path]
 572   if skip_mtab:
 573     cmd += ['-n']
 574   if fs_type:
 575     cmd += ['-t', fs_type]
 576   runcmd(cmd + ['-o', ','.join(mount_opts)], **kwargs)
 577
 578
 579 def MountTmpfsDir(path, name='osutils.tmpfs', size='5G',
 580                   mount_opts=('nodev', 'noexec', 'nosuid'), **kwargs):
 581   """Mount a tmpfs at |path|
 582
 583   Args:
 584     path: Directory to mount the tmpfs.
 585     name: Friendly name to include in mount output.
 586     size: Size of the temp fs.
 587     mount_opts: List of options to pass to `mount`.
 588     kwargs: Pass all other args to MountDir.
 589   """
 590   mount_opts = list(mount_opts) + ['size=%s' % size]
 591   MountDir(name, path, fs_type='tmpfs', mount_opts=mount_opts, **kwargs)
 592
 593
 594 def UmountDir(path, lazy=True, sudo=True, cleanup=True):
 595   """Unmount a previously mounted temp fs mount.
 596
 597   Args:
 598     path: Directory to unmount.
 599     lazy: Whether to do a lazy unmount.
 600     sudo: Run through sudo.
 601     cleanup: Whether to delete the |path| after unmounting.
 602              Note: Does not work when |lazy| is set.
 603   """
 604   if sudo:
 605     runcmd = cros_build_lib.SudoRunCommand
 606   else:
 607     runcmd = cros_build_lib.RunCommand
 608
 609   cmd = ['umount', '-d', path]
 610   if lazy:
 611     cmd += ['-l']
 612   runcmd(cmd)
 613
 614   if cleanup:
 615     # We will randomly get EBUSY here even when the umount worked.  Suspect
 616     # this is due to the host distro doing stupid crap on us like autoscanning
 617     # directories when they get mounted.
 618     def _retry(e):
 619       # When we're using `rm` (which is required for sudo), we can't cleanly
 620       # detect the aforementioned failure.  This is because `rm` will see the
 621       # errno, handle itself, and then do exit(1).  Which means all we see is
 622       # that rm failed.  Assume it's this issue as -rf will ignore most things.
 623       if isinstance(e, cros_build_lib.RunCommandError):
 624         return True
 625       else:
 626         # When we aren't using sudo, we do the unlink ourselves, so the exact
 627         # errno is bubbled up to us and we can detect it specifically without
 628         # potentially ignoring all other possible failures.
 629         return e.errno == errno.EBUSY
 630     retry_util.GenericRetry(_retry, 30, RmDir, path, sudo=sudo, sleep=60)
 631
 632
 633 def SetEnvironment(env):
 634   """Restore the environment variables to that of passed in dictionary."""
 635   os.environ.clear()
 636   os.environ.update(env)
 637
 638
 639 def SourceEnvironment(script, whitelist, ifs=',', env=None, multiline=False):
 640   """Returns the environment exported by a shell script.
 641
 642   Note that the script is actually executed (sourced), so do not use this on
 643   files that have side effects (such as modify the file system).  Stdout will
 644   be sent to /dev/null, so just echoing is OK.
 645
 646   Args:
 647     script: The shell script to 'source'.
 648     whitelist: An iterable of environment variables to retrieve values for.
 649     ifs: When showing arrays, what separator to use.
 650     env: A dict of the initial env to pass down.  You can also pass it None
 651          (to clear the env) or True (to preserve the current env).
 652     multiline: Allow a variable to span multiple lines.
 653
 654   Returns:
 655     A dictionary containing the values of the whitelisted environment
 656     variables that are set.
 657   """
 658   dump_script = ['source "%s" >/dev/null' % script,
 659                  'IFS="%s"' % ifs]
 660   for var in whitelist:
 661     dump_script.append(
 662         '[[ "${%(var)s+set}" == "set" ]] && echo %(var)s="${%(var)s[*]}"'
 663         % {'var': var})
 664   dump_script.append('exit 0')
 665
 666   if env is None:
 667     env = {}
 668   elif env is True:
 669     env = None
 670   output = cros_build_lib.RunCommand(['bash'], env=env, redirect_stdout=True,
 671                                      redirect_stderr=True, print_cmd=False,
 672                                      input='\n'.join(dump_script)).output
 673   return cros_build_lib.LoadKeyValueFile(cStringIO.StringIO(output),
 674                                          multiline=multiline)
 675
 676
 677 def ListBlockDevices(device_path=None, in_bytes=False):
 678   """Lists all block devices.
 679
 680   Args:
 681     device_path: device path (e.g. /dev/sdc).
 682     in_bytes: whether to display size in bytes.
 683
 684   Returns:
 685     A list of BlockDevice items with attributes 'NAME', 'RM', 'TYPE',
 686     'SIZE' (RM stands for removable).
 687   """
 688   keys = ['NAME', 'RM', 'TYPE', 'SIZE']
 689   BlockDevice = collections.namedtuple('BlockDevice', keys)
 690
 691   cmd = ['lsblk', '--pairs']
 692   if in_bytes:
 693     cmd.append('--bytes')
 694
 695   if device_path:
 696     cmd.append(device_path)
 697
 698   cmd += ['--output', ','.join(keys)]
 699   output = cros_build_lib.RunCommand(cmd, capture_output=True).output.strip()
 700   devices = []
 701   for line in output.splitlines():
 702     d = {}
 703     for k, v in re.findall(r'(\S+?)=\"(.+?)\"', line):
 704       d[k] = v
 705
 706     devices.append(BlockDevice(**d))
 707
 708   return devices
 709
 710
 711 def GetDeviceInfo(device, keyword='model'):
 712   """Get information of |device| by searching through device path.
 713
 714     Looks for the file named |keyword| in the path upwards from
 715     /sys/block/|device|/device. This path is a symlink and will be fully
 716     expanded when searching.
 717
 718   Args:
 719     device: Device name (e.g. 'sdc').
 720     keyword: The filename to look for (e.g. product, model).
 721
 722   Returns:
 723     The content of the |keyword| file.
 724   """
 725   device_path = os.path.join('/sys', 'block', device)
 726   if not os.path.isdir(device_path):
 727     raise ValueError('%s is not a valid device path.' % device_path)
 728
 729   path_list = ExpandPath(os.path.join(device_path, 'device')).split(os.path.sep)
 730   while len(path_list) > 2:
 731     target = os.path.join(os.path.sep.join(path_list), keyword)
 732     if os.path.isfile(target):
 733       return ReadFile(target).strip()
 734
 735     path_list = path_list[:-1]
 736
 737
 738 def GetDeviceSize(device_path, in_bytes=False):
 739   """Returns the size of |device|.
 740
 741   Args:
 742     device_path: Device path (e.g. '/dev/sdc').
 743     in_bytes: If set True, returns the size in bytes.
 744
 745   Returns:
 746     Size of the device in human readable format unless |in_bytes| is set.
 747   """
 748   devices = ListBlockDevices(device_path=device_path, in_bytes=in_bytes)
 749   for d in devices:
 750     if d.TYPE == 'disk':
 751       return int(d.SIZE) if in_bytes else d.SIZE
 752
 753   raise ValueError('No size info of %s is found.' % device_path)
 754
 755
 756 def GetExitStatus(status):
 757   """Get the exit status of a child from an os.waitpid call.
 758
 759   Args:
 760     status: The return value of os.waitpid(pid, 0)[1]
 761
 762   Returns:
 763     The exit status of the process. If the process exited with a signal,
 764     the return value will be 128 plus the signal number.
 765   """
 766   if os.WIFSIGNALED(status):
 767     return 128 + os.WTERMSIG(status)
 768   else:
 769     assert os.WIFEXITED(status), 'Unexpected exit status %r' % status
 770     return os.WEXITSTATUS(status)
 771
 772
 773 FileInfo = collections.namedtuple(
 774     'FileInfo', ['path', 'owner', 'size', 'atime', 'mtime'])
 775
 776
 777 def StatFilesInDirectory(path, recursive=False, to_string=False):
 778   """Stat files in the directory |path|.
 779
 780   Args:
 781     path: Path to the target directory.
 782     recursive: Whether to recurisvely list all files in |path|.
 783     to_string: Whether to return a string containing the metadata of the
 784       files.
 785
 786   Returns:
 787     If |to_string| is False, returns a list of FileInfo objects. Otherwise,
 788     returns a string of metadata of the files.
 789   """
 790   path = ExpandPath(path)
 791   def ToFileInfo(path, stat):
 792     return FileInfo(path,
 793                     pwd.getpwuid(stat.st_uid)[0],
 794                     stat.st_size,
 795                     datetime.datetime.fromtimestamp(stat.st_atime),
 796                     datetime.datetime.fromtimestamp(stat.st_mtime))
 797
 798   file_infos = []
 799   for root, dirs, files in os.walk(path, topdown=True):
 800     for filename in dirs + files:
 801       filepath = os.path.join(root, filename)
 802       file_infos.append(ToFileInfo(filepath, os.lstat(filepath)))
 803
 804     if not recursive:
 805       # Process only the top-most directory.
 806       break
 807
 808   if not to_string:
 809     return file_infos
 810
 811   msg = 'Listing the content of %s' % path
 812   msg_format = ('Path: {x.path}, Owner: {x.owner}, Size: {x.size} bytes, '
 813                 'Accessed: {x.atime}, Modified: {x.mtime}')
 814   msg = '%s\n%s' % (msg,
 815                     '\n'.join([msg_format.format(x=x) for x in file_infos]))
 816   return msg
 817
 818
 819 def MountImagePartition(image_file, part_number, destination, gpt_table=None,
 820                         sudo=True, makedirs=True, mount_opts=('ro', ),
 821                         skip_mtab=False):
 822   """Mount a |partition| from |image_file| to |destination|.
 823
 824   If there is a GPT table (GetImageDiskPartitionInfo), it will be used for
 825   start offset and size of the selected partition. Otherwise, the GPT will
 826   be read again from |image_file|. The GPT table MUST have unit of "B".
 827
 828   The mount option will be:
 829
 830     -o offset=XXX,sizelimit=YYY,(*mount_opts)
 831
 832   Args:
 833     image_file: A path to the image file (chromiumos_base_image.bin).
 834     part_number: A partition number.
 835     destination: A path to the mount point.
 836     gpt_table: A dictionary of PartitionInfo objects. See
 837       cros_build_lib.GetImageDiskPartitionInfo.
 838     sudo: Same as MountDir.
 839     makedirs: Same as MountDir.
 840     mount_opts: Same as MountDir.
 841     skip_mtab: Same as MountDir.
 842   """
 843
 844   if gpt_table is None:
 845     gpt_table = cros_build_lib.GetImageDiskPartitionInfo(image_file, 'B',
 846                                                          key_selector='number')
 847
 848   for _, part in gpt_table.items():
 849     if part.number == part_number:
 850       break
 851   else:
 852     part = None
 853     raise ValueError('Partition number %d not found in the GPT %r.' %
 854                      (part_number, gpt_table))
 855
 856   opts = ['loop', 'offset=%d' % part.start, 'sizelimit=%d' % part.size]
 857   opts += mount_opts
 858   MountDir(image_file, destination, sudo=sudo, makedirs=makedirs,
 859            mount_opts=opts, skip_mtab=skip_mtab)
 860
 861
 862 @contextlib.contextmanager
 863 def ChdirContext(target_dir):
 864   """A context manager to chdir() into |target_dir| and back out on exit.
 865
 866   Args:
 867     target_dir: A target directory to chdir into.
 868   """
 869
 870   cwd = os.getcwd()
 871   os.chdir(target_dir)
 872   try:
 873     yield
 874   finally:
 875     os.chdir(cwd)
 876
 877
 878 class MountImageContext(object):
 879   """A context manager to mount an image."""
 880
 881   def __init__(self, image_file, destination, part_selects=(1, 3)):
 882     """Construct a context manager object to actually do the job.
 883
 884     Specified partitions will be mounted under |destination| according to the
 885     pattern:
 886
 887       partition ---mount--> dir-<partition number>
 888
 889     Symlinks with labels "dir-<label>" will also be created in |destination| to
 890     point to the mounted partitions. If there is a conflict in symlinks, the
 891     first one wins.
 892
 893     The image is unmounted when this context manager exits.
 894
 895       with MountImageContext('build/images/wolf/latest', 'root_mount_point'):
 896         # "dir-1", and "dir-3" will be mounted in root_mount_point
 897         ...
 898
 899     Args:
 900       image_file: A path to the image file.
 901       destination: A directory in which all mount points and symlinks will be
 902         created. This parameter is relative to the CWD at the time __init__ is
 903         called.
 904       part_selects: A list of partition numbers or labels to be mounted. If an
 905         element is an integer, it is matched as partition number, otherwise
 906         a partition label.
 907     """
 908     self._image_file = image_file
 909     self._gpt_table = cros_build_lib.GetImageDiskPartitionInfo(
 910         self._image_file, 'B', key_selector='number'
 911     )
 912     # Target dir is absolute path so that we do not have to worry about
 913     # CWD being changed later.
 914     self._target_dir = ExpandPath(destination)
 915     self._part_selects = part_selects
 916     self._mounted = set()
 917     self._linked_labels = set()
 918
 919   def _GetMountPointAndSymlink(self, part):
 920     """Given a PartitionInfo, return a tuple of mount point and symlink.
 921
 922     Args:
 923       part: A PartitionInfo object.
 924
 925     Returns:
 926       A tuple (mount_point, symlink).
 927     """
 928     dest_number = os.path.join(self._target_dir, 'dir-%d' % part.number)
 929     dest_label = os.path.join(self._target_dir, 'dir-%s' % part.name)
 930     return dest_number, dest_label
 931
 932   def _Mount(self, part):
 933     """Mount the partition and create a symlink to the mount point.
 934
 935     The partition is mounted as "dir-partNumber", and the symlink "dir-label".
 936     If "dir-label" already exists, no symlink is created.
 937
 938     Args:
 939       part: A PartitionInfo object.
 940
 941     Raises:
 942       ValueError if mount point already exists.
 943     """
 944     if part in self._mounted:
 945       return
 946
 947     dest_number, dest_label = self._GetMountPointAndSymlink(part)
 948     if os.path.exists(dest_number):
 949       raise ValueError('Mount point %s already exists.' % dest_number)
 950
 951     MountImagePartition(self._image_file, part.number,
 952                         dest_number, self._gpt_table)
 953     self._mounted.add(part)
 954
 955     if not os.path.exists(dest_label):
 956       os.symlink(os.path.basename(dest_number), dest_label)
 957       self._linked_labels.add(dest_label)
 958
 959   def _Unmount(self, part):
 960     """Unmount a partition that was mounted by _Mount."""
 961     dest_number, dest_label = self._GetMountPointAndSymlink(part)
 962     # Due to crosbug/358933, the RmDir call might fail. So we skip the cleanup.
 963     UmountDir(dest_number, cleanup=False)
 964     self._mounted.remove(part)
 965
 966     if dest_label in self._linked_labels:
 967       SafeUnlink(dest_label)
 968       self._linked_labels.remove(dest_label)
 969
 970   def _CleanUp(self):
 971     """Unmount all mounted partitions."""
 972     to_be_rmdir = []
 973     for part in list(self._mounted):
 974       self._Unmount(part)
 975       dest_number, _ = self._GetMountPointAndSymlink(part)
 976       to_be_rmdir.append(dest_number)
 977     # Because _Unmount did not RmDir the mount points, we do that here.
 978     for path in to_be_rmdir:
 979       retry_util.RetryException(cros_build_lib.RunCommandError, 30,
 980                                 RmDir, path, sudo=True, sleep=60)
 981
 982   def __enter__(self):
 983     for selector in self._part_selects:
 984       matcher = operator.attrgetter('number')
 985       if not isinstance(selector, int):
 986         matcher = operator.attrgetter('name')
 987       for _, part in self._gpt_table.items():
 988         if matcher(part) == selector:
 989           try:
 990             self._Mount(part)
 991           except:
 992             self._CleanUp()
 993             raise
 994           break
 995       else:
 996         self._CleanUp()
 997         raise ValueError('Partition %r not found in the GPT %r.' %
 998                          (selector, self._gpt_table))
 999
1000     return self
1001
1002   def __exit__(self, exc_type, exc_value, traceback):
1003     self._CleanUp()
1004
1005
1006 MountInfo = collections.namedtuple('MountInfo',
1007     'source destination filesystem options')
1008
1009
1010 def IterateMountPoints(proc_file='/proc/mounts'):
1011   """Iterate over all mounts as reported by "/proc/mounts".
1012
1013   Args:
1014     proc_file: A path to a file whose content is similar to /proc/mounts.
1015       Default to "/proc/mounts" itself.
1016
1017   Returns:
1018     A generator that yields MountInfo objects.
1019   """
1020   with open(proc_file, 'rt') as f:
1021     for line in f:
1022       # Escape any \xxx to a char.
1023       source, destination, filesystem, options, _, _ = [
1024           re.sub(r'\\([0-7]{3})', lambda m: chr(int(m.group(1), 8)), x)
1025           for x in line.split()
1026       ]
1027       mtab = MountInfo(source, destination, filesystem, options)
1028       yield mtab