bisect_kit/arc_util.py

# -*- coding: utf-8 -*-
# Copyright 2018 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""ChromeOS ARC utility."""

from __future__ import print_function
import logging
import os
import re
import shutil
import subprocess
import tempfile

from bisect_kit import android_util
from bisect_kit import cros_util
from bisect_kit import util

logger = logging.getLogger(__name__)


def search_push_to_device_script_path(android_root):
  """Search push_to_device.py.

  Args:
    android_root: root path of android tree to search.

  Returns:
    the path of push_to_device.py; None if not found
  """
  # push_to_device.py has ever moved.
  for path in [
      'device/google/cheets2/scripts/push_to_device.py',
      'tools/vendor/google_prebuilts/arc/push_to_device.py',
  ]:
    fullpath = os.path.join(android_root, path)
    if os.path.exists(fullpath) and os.access(fullpath, os.X_OK):
      return path

  return None


def _push_to_device(dut, push_script, args, android_root=None, flavor=None):
  """Helper function to run push_to_device.py"""
  boot_id = cros_util.query_dut_boot_id(dut)

  assert push_script
  cmd = [push_script] + args
  stderr_lines = []

  def collect_stderr(line):
    stderr_lines.append(line)

  if android_root:
    assert flavor
    # New version of push_to_device can run without lunch, but go with lunch is
    # easier for now and compatible with old versions.
    android_util.lunch(
        android_root, flavor, *cmd, stderr_callback=collect_stderr)
  else:
    util.check_call(*cmd, stderr_callback=collect_stderr)

  stderr = ''.join(stderr_lines)
  if ('*** Reboot required. ***' in stderr or
      '*** Rebooting device. ***' in stderr or
      'root:Rebooting device...' in stderr):
    cros_util.wait_reboot_done(dut, boot_id)


def push_prebuilt_to_device(dut, flavor, build_id):
  """Pushes prebuilt ARC container to DUT.

  Args:
    dut: address of DUT
    flavor: android build flavor
    build_id: android build id.
  """
  # assume ARC's flavor is always product-variant
  assert '-' in flavor
  product, _ = flavor.split('-')
  build_filename = '{product}-img-{build_id}.zip'.format(
      product=product, build_id=build_id)
  sepolicy_filename = 'sepolicy.zip'

  tmp_dir = tempfile.mkdtemp()
  try:
    logger.debug('fetch and unpack push_to_device.zip to %s', tmp_dir)
    android_util.fetch_artifact(flavor, build_id, 'push_to_device.zip', tmp_dir)
    android_util.fetch_artifact(flavor, build_id, build_filename, tmp_dir)
    android_util.fetch_artifact(flavor, build_id, sepolicy_filename, tmp_dir)

    util.check_call('unzip', 'push_to_device.zip', cwd=tmp_dir)
    push_script = os.path.join(tmp_dir, 'push_to_device.py')

    # Don't use --use-prebuilt because push_to_device.py doesn't support
    # fetching artifacts with service account.
    args = [
        '--force',
        dut,
        '--use-prebuilt-file',
        os.path.join(tmp_dir, build_filename),
        '--sepolicy-artifacts-path',
        os.path.join(tmp_dir, sepolicy_filename),
    ]
    _push_to_device(dut, push_script, args)
  finally:
    logger.debug('clean up %s', tmp_dir)
    shutil.rmtree(tmp_dir)


def push_localbuild_to_device(android_root, dut, flavor):
  """Pushes localbuild ARC container to DUT.

  Args:
    android_root: root path of android tree.
    dut: address of DUT
    flavor: android build flavor
  """
  push_script = search_push_to_device_script_path(android_root)

  # Build push_to_device's dependency. We have to build it explicitly because
  # it's possible that no build is generated yet.
  # push_to_device.zip is valid since 2018-06.
  try:
    android_util.lunch(android_root, flavor, 'make', 'push_to_device.zip')
  except subprocess.CalledProcessError:
    # TODO(kcwu): remove this and build only push_to_device.zip once we don't
    #             care 2018-06 and earlier versions.
    android_util.lunch(android_root, flavor, 'make', 'unsquashfs',
                       'simg2img_host', 'mksquashfs', 'secilc')

  _push_to_device(dut, push_script, ['--force', dut], flavor)


def get_build_prop(host):
  """Get build property of a host.

  Args:
    host: host address

  Returns:
    dict of property
  """
  data = util.ssh_cmd(
      host, 'cat',
      '/opt/google/containers/android/rootfs/root/system/build.prop')
  if not data:
    return None
  result = {}
  for k, v in re.findall(r'^([^#=]+)=(.*)$', data, re.M):
    result[k] = v
  return result


def query_variant(host):
  """Queries DUT's build variant (ro.build.type)."""
  prop = get_build_prop(host)
  return prop and prop.get('ro.build.type')


def query_product(host):
  """Queries DUT's build product (ro.build.product)."""
  prop = get_build_prop(host)
  return prop and prop.get('ro.build.product')


def query_flavor(host):
  """Queries DUT's build flavor (ro.build.flavor)."""
  prop = get_build_prop(host)
  return prop and prop.get('ro.build.flavor')