#!/usr/bin/env python
#
# Copyright (c) 2013 The Chromium Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""Provisions Android devices with settings required for bots.

Usage:
  ./provision_devices.py [-d <device serial number>]
"""

import argparse
import datetime
import json
import logging
import os
import posixpath
import re
import subprocess
import sys
import time

from devil.android import battery_utils
from devil.android import device_blacklist
from devil.android import device_errors
from devil.android import device_temp_file
from devil.android import device_utils
from devil.android.sdk import version_codes
from devil.utils import run_tests_helper
from devil.utils import timeout_retry
from pylib import constants
from pylib import device_settings

_SYSTEM_WEBVIEW_PATHS = ['/system/app/webview', '/system/app/WebViewGoogle']
_CHROME_PACKAGE_REGEX = re.compile('.*chrom.*')
_TOMBSTONE_REGEX = re.compile('tombstone.*')


class _DEFAULT_TIMEOUTS(object):
  # L can take a while to reboot after a wipe.
  LOLLIPOP = 600
  PRE_LOLLIPOP = 180

  HELP_TEXT = '{}s on L, {}s on pre-L'.format(LOLLIPOP, PRE_LOLLIPOP)


class _PHASES(object):
  WIPE = 'wipe'
  PROPERTIES = 'properties'
  FINISH = 'finish'

  ALL = [WIPE, PROPERTIES, FINISH]


def ProvisionDevices(args):
  blacklist = (device_blacklist.Blacklist(args.blacklist_file)
               if args.blacklist_file
               else None)

  devices = device_utils.DeviceUtils.HealthyDevices(blacklist)
  if args.device:
    devices = [d for d in devices if d == args.device]
    if not devices:
      raise device_errors.DeviceUnreachableError(args.device)

  parallel_devices = device_utils.DeviceUtils.parallel(devices)
  parallel_devices.pMap(ProvisionDevice, blacklist, args)
  if args.auto_reconnect:
    _LaunchHostHeartbeat()
  blacklisted_devices = blacklist.Read() if blacklist else []
  if args.output_device_blacklist:
    with open(args.output_device_blacklist, 'w') as f:
      json.dump(blacklisted_devices, f)
  if all(d in blacklisted_devices for d in devices):
    raise device_errors.NoDevicesError
  return 0


def ProvisionDevice(device, blacklist, options):
  if options.reboot_timeout:
    reboot_timeout = options.reboot_timeout
  elif device.build_version_sdk >= version_codes.LOLLIPOP:
    reboot_timeout = _DEFAULT_TIMEOUTS.LOLLIPOP
  else:
    reboot_timeout = _DEFAULT_TIMEOUTS.PRE_LOLLIPOP

  def should_run_phase(phase_name):
    return not options.phases or phase_name in options.phases

  def run_phase(phase_func, reboot=True):
    try:
      device.WaitUntilFullyBooted(timeout=reboot_timeout, retries=0)
    except device_errors.CommandTimeoutError:
      logging.error('Device did not finish booting. Will try to reboot.')
      device.Reboot(timeout=reboot_timeout)
    phase_func(device, options)
    if reboot:
      device.Reboot(False, retries=0)
      device.adb.WaitForDevice()

  try:
    CheckExternalStorage(device)
    if should_run_phase(_PHASES.WIPE):
      if options.chrome_specific_wipe:
        run_phase(WipeChromeData)
      else:
        run_phase(WipeDevice)

    if should_run_phase(_PHASES.PROPERTIES):
      run_phase(SetProperties)

    if should_run_phase(_PHASES.FINISH):
      run_phase(FinishProvisioning, reboot=False)

    if options.chrome_specific_wipe:
      package = "com.google.android.gms"
      version_name = device.GetApplicationVersion(package)
      logging.info("Version name for %s is %s", package, version_name)

  except device_errors.CommandTimeoutError:
    logging.exception('Timed out waiting for device %s. Adding to blacklist.',
                      str(device))
    blacklist.Extend([str(device)])

  except device_errors.CommandFailedError:
    logging.exception('Failed to provision device %s. Adding to blacklist.',
                      str(device))
    blacklist.Extend([str(device)])

def CheckExternalStorage(device):
  """Checks that storage is writable and if not makes it writable.

  Arguments:
    device: The device to check.
  """
  try:
    with device_temp_file.DeviceTempFile(
        device.adb, suffix='.sh', dir=device.GetExternalStoragePath()):
      pass
  except device_errors.CommandFailedError:
    logging.info('External storage not writable. Remounting / as RW')
    device.RunShellCommand(['mount', '-o', 'remount,rw', '/'],
                           check_return=True, as_root=True)
    with device_temp_file.DeviceTempFile(
        device.adb, suffix='.sh', dir=device.GetExternalStoragePath()):
      pass

def WipeChromeData(device, options):
  """Wipes chrome specific data from device

  (1) uninstall any app whose name matches *chrom*, except
      com.android.chrome, which is the chrome stable package. Doing so also
      removes the corresponding dirs under /data/data/ and /data/app/
  (2) remove any dir under /data/app-lib/ whose name matches *chrom*
  (3) remove any files under /data/tombstones/ whose name matches "tombstone*"
  (4) remove /data/local.prop if there is any
  (5) remove /data/local/chrome-command-line if there is any
  (6) remove anything under /data/local/.config/ if the dir exists
      (this is telemetry related)
  (7) remove anything under /data/local/tmp/

  Arguments:
    device: the device to wipe
  """
  if options.skip_wipe:
    return

  try:
    device.EnableRoot()
    _UninstallIfMatch(device, _CHROME_PACKAGE_REGEX,
                      constants.PACKAGE_INFO['chrome_stable'].package)
    _WipeUnderDirIfMatch(device, '/data/app-lib/', _CHROME_PACKAGE_REGEX)
    _WipeUnderDirIfMatch(device, '/data/tombstones/', _TOMBSTONE_REGEX)

    _WipeFileOrDir(device, '/data/local.prop')
    _WipeFileOrDir(device, '/data/local/chrome-command-line')
    _WipeFileOrDir(device, '/data/local/.config/')
    _WipeFileOrDir(device, '/data/local/tmp/')

    device.RunShellCommand('rm -rf %s/*' % device.GetExternalStoragePath(),
                           check_return=True)
  except device_errors.CommandFailedError:
    logging.exception('Possible failure while wiping the device. '
                      'Attempting to continue.')


def WipeDevice(device, options):
  """Wipes data from device, keeping only the adb_keys for authorization.

  After wiping data on a device that has been authorized, adb can still
  communicate with the device, but after reboot the device will need to be
  re-authorized because the adb keys file is stored in /data/misc/adb/.
  Thus, adb_keys file is rewritten so the device does not need to be
  re-authorized.

  Arguments:
    device: the device to wipe
  """
  if options.skip_wipe:
    return

  try:
    device.EnableRoot()
    device_authorized = device.FileExists(constants.ADB_KEYS_FILE)
    if device_authorized:
      adb_keys = device.ReadFile(constants.ADB_KEYS_FILE,
                                 as_root=True).splitlines()
    device.RunShellCommand(['wipe', 'data'],
                           as_root=True, check_return=True)
    device.adb.WaitForDevice()

    if device_authorized:
      adb_keys_set = set(adb_keys)
      for adb_key_file in options.adb_key_files or []:
        try:
          with open(adb_key_file, 'r') as f:
            adb_public_keys = f.readlines()
          adb_keys_set.update(adb_public_keys)
        except IOError:
          logging.warning('Unable to find adb keys file %s.', adb_key_file)
      _WriteAdbKeysFile(device, '\n'.join(adb_keys_set))
  except device_errors.CommandFailedError:
    logging.exception('Possible failure while wiping the device. '
                      'Attempting to continue.')


def _WriteAdbKeysFile(device, adb_keys_string):
  dir_path = posixpath.dirname(constants.ADB_KEYS_FILE)
  device.RunShellCommand(['mkdir', '-p', dir_path],
                         as_root=True, check_return=True)
  device.RunShellCommand(['restorecon', dir_path],
                         as_root=True, check_return=True)
  device.WriteFile(constants.ADB_KEYS_FILE, adb_keys_string, as_root=True)
  device.RunShellCommand(['restorecon', constants.ADB_KEYS_FILE],
                         as_root=True, check_return=True)


def SetProperties(device, options):
  try:
    device.EnableRoot()
  except device_errors.CommandFailedError as e:
    logging.warning(str(e))

  _ConfigureLocalProperties(device, options.enable_java_debug)
  device_settings.ConfigureContentSettings(
      device, device_settings.DETERMINISTIC_DEVICE_SETTINGS)
  if options.disable_location:
    device_settings.ConfigureContentSettings(
        device, device_settings.DISABLE_LOCATION_SETTINGS)
  else:
    device_settings.ConfigureContentSettings(
        device, device_settings.ENABLE_LOCATION_SETTINGS)

  if options.disable_mock_location:
    device_settings.ConfigureContentSettings(
        device, device_settings.DISABLE_MOCK_LOCATION_SETTINGS)
  else:
    device_settings.ConfigureContentSettings(
        device, device_settings.ENABLE_MOCK_LOCATION_SETTINGS)

  device_settings.SetLockScreenSettings(device)
  if options.disable_network:
    device_settings.ConfigureContentSettings(
        device, device_settings.NETWORK_DISABLED_SETTINGS)

  if options.disable_system_chrome:
    # The system chrome version on the device interferes with some tests.
    device.RunShellCommand(['pm', 'disable', 'com.android.chrome'],
                           check_return=True)

  if options.remove_system_webview:
    if device.HasRoot():
      # This is required, e.g., to replace the system webview on a device.
      device.adb.Remount()
      device.RunShellCommand(['stop'], check_return=True)
      device.RunShellCommand(['rm', '-rf'] + _SYSTEM_WEBVIEW_PATHS,
                             check_return=True)
      device.RunShellCommand(['start'], check_return=True)
    else:
      logging.warning('Cannot remove system webview from a non-rooted device')


def _ConfigureLocalProperties(device, java_debug=True):
  """Set standard readonly testing device properties prior to reboot."""
  local_props = [
      'persist.sys.usb.config=adb',
      'ro.monkey=1',
      'ro.test_harness=1',
      'ro.audio.silent=1',
      'ro.setupwizard.mode=DISABLED',