gooftool: move class Util to common.py

Move utility class Util to common.py to make it easiler to reuse by
other modules.

BUG=None
TEST=make test

Change-Id: I6d13dadeb361c43ab9e20fa07f71c20bae362dd4
Reviewed-on: https://chromium-review.googlesource.com/336376
Commit-Ready: Wei-Han Chen <stimim@chromium.org>
Tested-by: Wei-Han Chen <stimim@chromium.org>
Reviewed-by: Hung-Te Lin <hungte@chromium.org>
diff --git a/py/gooftool/common.py b/py/gooftool/common.py
index 8268497..bfde42e 100644
--- a/py/gooftool/common.py
+++ b/py/gooftool/common.py
@@ -6,10 +6,14 @@
 
 from __future__ import print_function
 
+import os
+import re
 import logging
 from subprocess import Popen, PIPE
 
 import factory_common  # pylint: disable=W0611
+from cros.factory.utils import sys_utils
+from cros.factory.utils.type_utils import Error
 from cros.factory.utils.type_utils import Obj
 
 
@@ -33,3 +37,279 @@
                   (', stderr: %s' % repr(stderr.strip()) if stderr else ''))
   status = process.poll()
   return Obj(stdout=stdout, stderr=stderr, status=status, success=(status == 0))
+
+
+class Util(object):
+  """A collection of util functions that Gooftool needs."""
+
+  def __init__(self):
+    self.shell = Shell
+
+  def _IsDeviceFixed(self, dev):
+    """Check if a device is a fixed device, i.e. not a removable device.
+
+    Args:
+      dev: A device string under /sys/block.
+
+    Returns:
+      True if the given device is fixed, and false if it is not.
+    """
+
+    sysfs_path = '/sys/block/%s/removable' % dev
+    return (os.path.exists(sysfs_path) and
+            open(sysfs_path).read().strip() == '0')
+
+  def GetPrimaryDevicePath(self, partition=None):
+    """Gets the path for the primary device, which is the only non-removable
+    device in the system.
+
+    Args:
+      partition: The index of the partition on primary device.
+
+    Returns:
+      The path to the primary device. If partition is specified, the path
+      points to that partition of the primary device. e.g. /dev/sda1
+    """
+
+    alpha_re = re.compile(r'^/dev/([a-zA-Z]+)[0-9]+$')
+    alnum_re = re.compile(r'^/dev/([a-zA-Z]+[0-9]+)p[0-9]+$')
+    matched_alnum = False
+    dev_set = set()
+    for path in self.shell('cgpt find -t rootfs').stdout.strip().split():
+      for dev in alpha_re.findall(path):
+        if self._IsDeviceFixed(dev):
+          dev_set.add(dev)
+          matched_alnum = False
+      for dev in alnum_re.findall(path):
+        if self._IsDeviceFixed(dev):
+          dev_set.add(dev)
+          matched_alnum = True
+    if len(dev_set) != 1:
+      raise Error('zero or multiple primary devs: %s' % dev_set)
+    dev_path = os.path.join('/dev', dev_set.pop())
+    if partition is None:
+      return dev_path
+    fmt_str = '%sp%d' if matched_alnum else '%s%d'
+    return fmt_str % (dev_path, partition)
+
+  def GetPartitionDevice(self, path):
+    """Returns a device path string from partition path.
+
+    /dev/sda1 => /dev/sda.
+    /dev/mmcblk0p2 => /dev/mmcblk0.
+    """
+    return ''.join(re.findall(
+        r'(.*[^0-9][0-9]+)p[0-9]+|(.*[^0-9])[0-9]+', path)[0])
+
+  def FindScript(self, script_name):
+    """Finds the script under /usr/local/factory/sh
+
+    Args:
+      script_name: The name of the script to look for.
+
+    Returns:
+      The path of the found script.
+
+    Raises:
+      Error if the script is not found.
+    """
+
+    # __file__ is in /usr/local/factory/py/gooftool/__init__.py
+    factory_base = os.path.realpath(os.path.join(
+        os.path.dirname(os.path.realpath(__file__)), '..', '..'))
+    script_path = os.path.join(factory_base, 'sh', script_name)
+    if not os.path.isfile(script_path):
+      raise Error('Needed script %s does not exist.' % script_path)
+    return script_path
+
+  def FindAndRunScript(self, script_name, post_opts=None, pre_opts=None):
+    """Finds and runs the script with given options.
+
+    Args:
+      script_name: The name of the script to look up and run.
+      post_opts: A list of strings that will be appended in the command after
+        the script's name.
+      pre_opts: A list of strings that will be prepended in the command before
+        the script's name.
+
+    Returns:
+      The result of execusion.
+
+    Raises:
+      Error if execusion failed.
+    """
+
+    assert not post_opts or isinstance(post_opts, list)
+    assert not pre_opts or isinstance(pre_opts, list)
+
+    script = self.FindScript(script_name)
+    cmd = '%s %s %s' % (' '.join(pre_opts) if pre_opts else '',
+                        script,
+                        ' '.join(post_opts) if post_opts else '')
+    result = self.shell(cmd.strip())
+    if not result.success:
+      raise Error, '%r failed, stderr: %r' % (cmd, result.stderr)
+
+    return result
+
+  def GetReleaseRootPartitionPath(self):
+    """Gets the path for release root partition."""
+
+    return self.GetPrimaryDevicePath(5)
+
+  def GetReleaseKernelPartitionPath(self):
+    """Gets the path for release kernel partition."""
+
+    return self.GetPrimaryDevicePath(4)
+
+  def GetReleaseImageLsbData(self):
+    """Gets the /etc/lsb-release content from release image partition.
+
+    Returns:
+      A dictionary containing the key-value pairs in lsb-release.
+    """
+    lsb_content = sys_utils.MountDeviceAndReadFile(
+        self.GetReleaseRootPartitionPath(), 'etc/lsb-release')
+    return dict(re.findall('^(.+)=(.+)$', lsb_content, re.MULTILINE))
+
+  def GetAllowedReleaseImageChannels(self):
+    """Returns a list of channels allowed for release image."""
+    return ['dev', 'beta', 'stable']
+
+  def GetReleaseImageChannel(self):
+    """Returns the channel of current release image."""
+    return self.GetReleaseImageLsbData().get('CHROMEOS_RELEASE_TRACK')
+
+  def GetReleaseImageVersion(self):
+    """Returns the current release image version."""
+    return self.GetReleaseImageLsbData().get('GOOGLE_RELEASE')
+
+  def GetVBSharedDataFlags(self):
+    """Gets VbSharedData flags.
+
+    Returns:
+      An integer representation of the flags.
+    """
+
+    return int(self.shell('crossystem vdat_flags').stdout.strip(), 0)
+
+  def GetCurrentDevSwitchPosition(self):
+    """Gets the position for the current developer switch.
+
+    Returns:
+      An integer representation of the current developer switch position.
+    """
+    return int(self.shell('crossystem devsw_cur').stdout.strip(), 0)
+
+  def GetCrosSystem(self):
+    """Gets the output of 'crossystem'.
+
+    Returns:
+      A dict for key-value pairs for the output of 'crossystem'.
+      e.g. {'flag_name': 'flag_value'}
+    """
+    crossystem_result = self.shell('crossystem').stdout.strip().splitlines()
+    # The crossytem output contains many lines like:
+    # 'key = value  # description'
+    # Use regexps to pull out the key-value pairs and build a dict.
+    # Note that value could also contain equal signs.
+    output = {}
+    for entry in crossystem_result:
+      # Any unrecognized format should fail here.
+      key, value = re.findall(r'\A(\S+)\s+=\s+(.*)#.*\Z', entry)[0]
+      output[key] = value.strip()
+
+    return output
+
+  def GetCgptAttributes(self, device=None):
+    if device is None:
+      device = self.GetPrimaryDevicePath()
+
+    attrs = {}
+    for line in self.shell('cgpt show %s -q' % device).stdout.splitlines():
+      # format: offset size no name
+      part_no = line.split()[2]
+      attrs[part_no] = self.shell('cgpt show %s -i %s -A' %
+                                  (device, part_no)).stdout.strip()
+    return attrs
+
+  def SetCgptAttributes(self, attrs, device=None):
+    if device is None:
+      device = self.GetPrimaryDevicePath()
+
+    curr_attrs = self.GetCgptAttributes()
+    for k, v in attrs.iteritems():
+      if curr_attrs.get(k) == v:
+        continue
+      if not self.shell('cgpt add %s -i %s -A %s' % (device, k, v)).success:
+        raise Error('Failed to set device config: %s#%s=%s' % (device, k, v))
+
+  def InvokeChromeOSPostInstall(self, root_dev=None):
+    """Invokes the ChromeOS post-install script (/postinst)."""
+    if root_dev is None:
+      root_dev = self._util.GetReleaseRootPartitionPath()
+
+    logging.info('Running ChromeOS post-install on %s...', root_dev)
+
+    # Some compatible and experimental fs (e.g., ext4) may be buggy and still
+    # try to write the file system even if we mount it with "ro" (ex, when
+    # seeing journaling error in ext3, or s_kbytes_written in ext4). It is
+    # safer to always mount the partition with legacy ext2. (ref:
+    # chrome-os-partner:3940)
+    with sys_utils.MountPartition(root_dev, fstype='ext2') as mount_path:
+      # IS_FACTORY_INSTALL is used to prevent postinst trying to update firmare.
+      command = ('IS_FACTORY_INSTALL=1 IS_INSTALL=1 "%s"/postinst %s' %
+                 (mount_path, root_dev))
+      result = self.shell(command)
+      if not result.success:
+        raise Error('chromeos-postinst on %s failed with error: code=%s. %s' %
+                    (root_dev, result.status, result.stderr))
+
+  def EnableKernel(self, device, part_no):
+    """Enables the kernel partition from GPT."""
+    logging.info('Enabling kernel on %s#%s...', device, part_no)
+    r = self.shell('cgpt add -i %s -P 3 -S 1 -T 0 %s' % (part_no, device))
+    if not r.success:
+      raise Error('Failed to enable kernel on %s#%s' % (device, part_no))
+
+  def DisableKernel(self, device, part_no):
+    """Disables the kernel partition from GPT."""
+    logging.info('Disabling kernel on %s#%s...', device, part_no)
+    r = self.shell('cgpt add -i %s -P 0 -S 0 -T 0 %s' % (part_no, device))
+    if not r.success:
+      raise Error('Failed to disable kernel on %s#%s' % (device, part_no))
+
+  def IsChromeOSFirmware(self):
+    """Returns if the system is running ChromeOS firmware."""
+    r = self.shell('crossystem mainfw_type')
+    return r.success and r.stdout.strip() != 'nonchrome'
+
+  def EnableReleasePartition(self, root_dev):
+    """Enables a release image partition on disk."""
+    # TODO(hungte) replce sh/enable_release_partition.sh
+    release_no = int(root_dev[-1]) - 1
+    factory_map = {2: 4, 4: 2}
+    if release_no not in factory_map:
+      raise ValueError('EnableReleasePartition: Cannot identify kernel %s' %
+                       root_dev)
+
+    factory_no = factory_map[release_no]
+    device = self.GetPartitionDevice(root_dev)
+    curr_attrs = self.GetCgptAttributes(device)
+    try:
+      # When booting with legacy firmware, we need to update the legacy boot
+      # loaders to activate new kernel; on a real ChromeOS firmware, only CGPT
+      # header is used, and postinst is already performed in verify_rootfs.
+      if self.IsChromeOSFirmware():
+        self.InvokeChromeOSPostInstall(root_dev)
+      self.shell('crossystem disable_dev_request=1')
+      self.DisableKernel(device, factory_no)
+      self.EnableKernel(device, release_no)
+      # Enforce a sync and wait for underlying hardware to flush.
+      logging.info('Syncing disks...')
+      self.shell('sync; sleep 3')
+      logging.info('Enable release partition: Complete.')
+    except:  # pylint: disable=bare-except
+      logging.error('FAIL: Failed to enable release partition.')
+      self.shell('crossystem disable_dev_request=0')
+      self.SetCgptAttributes(curr_attrs, device)