# -*- coding: utf-8 -*-
# Copyright (c) 2012 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.

"""Purpose of this module is to hold common script/commandline functionality.

This ranges from optparse, to a basic script wrapper setup (much like
what is used for chromite.bin.*).
"""

from __future__ import print_function

import argparse
import collections
import datetime
import functools
import os
import optparse  # pylint: disable=deprecated-module
import signal
import sys

import six
from six.moves import urllib

from autotest_lib.utils.frozen_chromite.lib import constants
from autotest_lib.utils.frozen_chromite.lib import cros_build_lib
from autotest_lib.utils.frozen_chromite.lib import cros_collections
from autotest_lib.utils.frozen_chromite.lib import cros_logging as logging
from autotest_lib.utils.frozen_chromite.lib import gs
from autotest_lib.utils.frozen_chromite.lib import osutils
from autotest_lib.utils.frozen_chromite.lib import path_util
from autotest_lib.utils.frozen_chromite.lib import terminal
from autotest_lib.utils.frozen_chromite.utils import attrs_freezer


DEVICE_SCHEME_FILE = 'file'
DEVICE_SCHEME_SERVO = 'servo'
DEVICE_SCHEME_SSH = 'ssh'
DEVICE_SCHEME_USB = 'usb'


class ChrootRequiredError(Exception):
  """Raised when a command must be run in the chroot

  This exception is intended to be caught by code which will restart execution
  in the chroot. Throwing this exception allows contexts to be exited and
  general cleanup to happen before we exec an external binary.

  The command to run inside the chroot, and (optionally) special cros_sdk
  arguments are attached to the exception. Any adjustments to the arguments
  should be done before raising the exception.
  """
  def __init__(self, cmd, chroot_args=None, extra_env=None):
    """Constructor for ChrootRequiredError.

    Args:
      cmd: Command line to run inside the chroot as a list of strings.
      chroot_args: Arguments to pass directly to cros_sdk.
      extra_env: Environmental variables to set in the chroot.
    """
    super(ChrootRequiredError, self).__init__()
    self.cmd = cmd
    self.chroot_args = chroot_args
    self.extra_env = extra_env


class ExecRequiredError(Exception):
  """Raised when a command needs to exec, after cleanup.

  This exception is intended to be caught by code which will exec another
  command. Throwing this exception allows contexts to be exited and general
  cleanup to happen before we exec an external binary.

  The command to run is attached to the exception. Any adjustments to the
  arguments should be done before raising the exception.
  """
  def __init__(self, cmd):
    """Constructor for ExecRequiredError.

    Args:
      cmd: Command line to run inside the chroot as a list of strings.
    """
    super(ExecRequiredError, self).__init__()
    self.cmd = cmd


def AbsolutePath(_option, _opt, value):
  """Expand paths and make them absolute."""
  return osutils.ExpandPath(value)


def NormalizeGSPath(value):
  """Normalize GS paths."""
  url = gs.CanonicalizeURL(value, strict=True)
  return '%s%s' % (gs.BASE_GS_URL, os.path.normpath(url[len(gs.BASE_GS_URL):]))


def NormalizeLocalOrGSPath(value):
  """Normalize a local or GS path."""
  ptype = 'gs_path' if gs.PathIsGs(value) else 'path'
  return VALID_TYPES[ptype](value)


def NormalizeAbUrl(value):
  """Normalize an androidbuild URL."""
  if not value.startswith('ab://'):
    # Give a helpful error message about the format expected.  Putting this
    # message in the exception is useless because argparse ignores the
    # exception message and just says the value is invalid.
    msg = 'Invalid ab:// URL format: [%s].' % value
    logging.error(msg)
    raise ValueError(msg)

  # If no errors, just return the unmodified value.
  return value


def ValidateCipdURL(value):
  """Return plain string."""
  if not value.startswith('cipd://'):
    msg = 'Invalid cipd:// URL format: %s' % value
    logging.error(msg)
    raise ValueError(msg)
  return value


def ParseBool(value):
  """Parse bool argument into a bool value.

  For the existing type=bool functionality, the parser uses the built-in bool(x)
  function to determine the value.  This function will only return false if x
  is False or omitted.  Even with this type specified, however, arguments that
  are generated from a command line initially get parsed as a string, and for
  any string value passed in to bool(x), it will always return True.

  Args:
    value: String representing a boolean value.

  Returns:
    True or False.
  """
  return cros_build_lib.BooleanShellValue(value, False)


def ParseDate(value):
  """Parse date argument into a datetime.date object.

  Args:
    value: String representing a single date in "YYYY-MM-DD" format.

  Returns:
    A datetime.date object.
  """
  try:
    return datetime.datetime.strptime(value, '%Y-%m-%d').date()
  except ValueError:
    # Give a helpful error message about the format expected.  Putting this
    # message in the exception is useless because argparse ignores the
    # exception message and just says the value is invalid.
    logging.error('Date is expected to be in format YYYY-MM-DD.')
    raise


def NormalizeUri(value):
  """Normalize a local path or URI."""
  o = urllib.parse.urlparse(value)
  if o.scheme == 'file':
    # Trim off the file:// prefix.
    return VALID_TYPES['path'](value[7:])
  elif o.scheme not in ('', 'gs'):
    o = list(o)
    o[2] = os.path.normpath(o[2])
    return urllib.parse.urlunparse(o)
  else:
    return NormalizeLocalOrGSPath(value)


# A Device object holds information parsed from the command line input:
#   scheme: DEVICE_SCHEME_SSH, DEVICE_SCHEME_USB, DEVICE_SCHEME_SERVO,
#     or DEVICE_SCHEME_FILE.
#   username: String SSH username or None.
#   hostname: String SSH hostname or None.
#   port: Int SSH or Servo port or None.
#   path: String USB/file path or None.
#   raw: String raw input from the command line.
#   serial_number: String Servo serial number or None.
# For now this is a superset of all information for USB, SSH, or file devices.
# If functionality diverges based on type, it may be useful to split this into
# separate device classes instead.
Device = cros_collections.Collection(
    'Device', scheme=None, username=None, hostname=None, port=None, path=None,
    raw=None, serial_number=None)


class DeviceParser(object):
  """Parses devices as an argparse argument type.

  In addition to parsing user input, this class will also ensure that only
  supported device schemes are accepted by the parser. For example,
  `cros deploy` only makes sense with an SSH device, but `cros flash` can use
  SSH, USB, or file device schemes.

  If the device input is malformed or the scheme is wrong, an error message will
  be printed and the program will exit.

  Valid device inputs are:
    - [ssh://][username@]hostname[:port].
    - usb://[path].
    - file://path or /absolute_path.
    - servo:port[:port] to use a port via dut-control, e.g. servo:port:1234.
    - servo:serial:serial-number to use the servo's serial number,
        e.g. servo:serial:641220-00057 servo:serial:C1230024192.
    - [ssh://]:vm:.

  The last item above is an alias for ssh'ing into a virtual machine on a
  localhost.  It gets translated into 'localhost:9222'.

  Examples:
    parser = argparse.ArgumentParser()

    parser.add_argument(
      'ssh_device',
      type=commandline.DeviceParser(commandline.DEVICE_SCHEME_SSH))

    parser.add_argument(
      'usb_or_file_device',
      type=commandline.DeviceParser([commandline.DEVICE_SCHEME_USB,
                                     commandline.DEVICE_SCHEME_FILE]))
  """

  def __init__(self, schemes):
    """Initializes the parser.

    See the class comments for usage examples.

    Args:
      schemes: A scheme or list of schemes to accept.
    """
    self.schemes = ([schemes] if isinstance(schemes, six.string_types)
                    else schemes)
    # Provide __name__ for argparse to print on failure, or else it will use
    # repr() which creates a confusing error message.
    self.__name__ = type(self).__name__

  def __call__(self, value):
    """Parses a device input and enforces constraints.

    DeviceParser is an object so that a set of valid schemes can be specified,
    but argparse expects a parsing function, so we overload __call__() for
    argparse to use.

    Args:
      value: String representing a device target. See class comments for
        valid device input formats.

    Returns:
      A Device object.

    Raises:
      ValueError: |value| is not a valid device specifier or doesn't
        match the supported list of schemes.
    """
    try:
      device = self._ParseDevice(value)
      self._EnforceConstraints(device, value)
      return device
    except ValueError as e:
      # argparse ignores exception messages, so print the message manually.
      logging.error(e)
      raise
    except Exception as e:
      logging.error('Internal error while parsing device input: %s', e)
      raise

  def _EnforceConstraints(self, device, value):
    """Verifies that user-specified constraints are upheld.

    Checks that the parsed device has a scheme that matches what the user
    expects. Additional constraints can be added if needed.

    Args:
      device: Device object.
      value: String representing a device target.

    Raises:
      ValueError: |device| has the wrong scheme.
    """
    if device.scheme not in self.schemes:
      raise ValueError('Unsupported scheme "%s" for device "%s"' %
                       (device.scheme, value))

  def _ParseDevice(self, value):
    """Parse a device argument.

    Args:
      value: String representing a device target.

    Returns:
      A Device object.

    Raises:
      ValueError: |value| is not a valid device specifier.
    """
    # ':vm:' is an alias for ssh'ing into a virtual machihne on localhost;
    # translate it appropriately.
    if value.strip().lower() == ':vm:':
      value = 'localhost:9222'
    elif value.strip().lower() == 'ssh://:vm:':
      value = 'ssh://localhost:9222'
    parsed = urllib.parse.urlparse(value)

    # crbug.com/1069325: Starting in python 3.7 urllib has different parsing
    # results. 127.0.0.1:9999 parses as scheme='127.0.0.1' path='9999'
    # instead of scheme='' path='127.0.0.1:9999'. We want that parsed as ssh.
    # Check for '.' or 'localhost' in the scheme to catch the most common cases
    # for this result.
    if (not parsed.scheme or '.' in parsed.scheme or
        parsed.scheme == 'localhost'):
      # Default to a file scheme for absolute paths, SSH scheme otherwise.
      if value and value[0] == '/':
        scheme = DEVICE_SCHEME_FILE
      else:
        # urlparse won't provide hostname/username/port unless a scheme is
        # specified so we need to re-parse.
        parsed = urllib.parse.urlparse('%s://%s' % (DEVICE_SCHEME_SSH, value))
        scheme = DEVICE_SCHEME_SSH
    else:
      scheme = parsed.scheme.lower()

    if scheme == DEVICE_SCHEME_SSH:
      hostname = parsed.hostname
      port = parsed.port
      if hostname == 'localhost' and not port:
        # Use of localhost as the actual machine is uncommon enough relative to
        # the use of KVM that we require users to specify localhost:22 if they
        # actually want to connect to the localhost.  Otherwise the expectation
        # is that they intend to access the VM but forget or didn't know to use
        # port 9222.
        raise ValueError('To connect to localhost, use ssh://localhost:22 '
                         'explicitly, or use ssh://localhost:9222 for the local'
                         ' VM.')
      if not hostname:
        raise ValueError('Hostname is required for device "%s"' % value)
      return Device(scheme=scheme, username=parsed.username, hostname=hostname,
                    port=port, raw=value)
    elif scheme == DEVICE_SCHEME_USB:
      path = parsed.netloc + parsed.path
      # Change path '' to None for consistency.
      return Device(scheme=scheme, path=path if path else None, raw=value)
    elif scheme == DEVICE_SCHEME_FILE:
      path = parsed.netloc + parsed.path
      if not path:
        raise ValueError('Path is required for "%s"' % value)
      return Device(scheme=scheme, path=path, raw=value)
    elif scheme == DEVICE_SCHEME_SERVO:
      # Parse the identifier type and value.
      servo_type, _, servo_id = parsed.path.partition(':')
      # Don't want to do the netloc before the split in case of serial number.
      servo_type = servo_type.lower()

      return self._parse_servo(servo_type, servo_id)
    else:
      raise ValueError('Unknown device scheme "%s" in "%s"' % (scheme, value))

  @staticmethod
  def _parse_servo(servo_type, servo_id):
    """Parse a servo device from the parsed servo uri info.

    Args:
      servo_type: The servo identifier type, either port or serial.
      servo_id: The servo identifier, either the port number it is
        communicating through or its serial number.
    """
    servo_port = None
    serial_number = None
    if servo_type == 'serial':
      if servo_id:
        serial_number = servo_id
      else:
        raise ValueError('No serial number given.')
    elif servo_type == 'port':
      if servo_id:
        # Parse and validate when given.
        try:
          servo_port = int(servo_id)
        except ValueError:
          raise ValueError('Invalid servo port value: %s' % servo_id)
        if servo_port <= 0 or servo_port > 65535:
          raise ValueError(
              'Invalid port, must be 1-65535: %d given.' % servo_port)
    else:
      raise ValueError('Invalid servo type given: %s' % servo_type)

    return Device(
        scheme=DEVICE_SCHEME_SERVO,
        port=servo_port,
        serial_number=serial_number)


class _AppendOption(argparse.Action):
  """Append the command line option (with no arguments) to dest.

  parser.add_argument('-b', '--barg', dest='out', action='append_option')
  options = parser.parse_args(['-b', '--barg'])
  options.out == ['-b', '--barg']
  """
  def __init__(self, option_strings, dest, **kwargs):
    if 'nargs' in kwargs:
      raise ValueError('nargs is not supported for append_option action')
    super(_AppendOption, self).__init__(
        option_strings, dest, nargs=0, **kwargs)

  def __call__(self, parser, namespace, values, option_string=None):
    if getattr(namespace, self.dest, None) is None:
      setattr(namespace, self.dest, [])
    getattr(namespace, self.dest).append(option_string)


class _AppendOptionValue(argparse.Action):
  """Append the command line option to dest. Useful for pass along arguments.

  parser.add_argument('-b', '--barg', dest='out', action='append_option_value')
  options = parser.parse_args(['--barg', 'foo', '-b', 'bar'])
  options.out == ['-barg', 'foo', '-b', 'bar']
  """
  def __call__(self, parser, namespace, values, option_string=None):
    if getattr(namespace, self.dest, None) is None:
      setattr(namespace, self.dest, [])
    getattr(namespace, self.dest).extend([option_string, str(values)])


class _SplitExtendAction(argparse.Action):
  """Callback to split the argument and extend existing value.

  We normalize whitespace before splitting.  This is to support the forms:
    cbuildbot -p 'proj:branch ' ...
    cbuildbot -p ' proj:branch' ...
    cbuildbot -p 'proj:branch  proj2:branch' ...
    cbuildbot -p "$(some_command_that_returns_nothing)" ...
  """
  def __call__(self, parser, namespace, values, option_string=None):
    if getattr(namespace, self.dest, None) is None:
      setattr(namespace, self.dest, [])
    getattr(namespace, self.dest).extend(values.split())


VALID_TYPES = {
    'ab_url': NormalizeAbUrl,
    'bool': ParseBool,
    'cipd': ValidateCipdURL,
    'date': ParseDate,
    'path': osutils.ExpandPath,
    'gs_path': NormalizeGSPath,
    'local_or_gs_path': NormalizeLocalOrGSPath,
    'path_or_uri': NormalizeUri,
}

VALID_ACTIONS = {
    'append_option': _AppendOption,
    'append_option_value': _AppendOptionValue,
    'split_extend': _SplitExtendAction,
}

_DEPRECATE_ACTIONS = [None, 'store', 'store_const', 'store_true', 'store_false',
                      'append', 'append_const', 'count'] + list(VALID_ACTIONS)


class _DeprecatedAction(object):
  """Base functionality to allow adding warnings for deprecated arguments.

  To add a deprecated warning, simply include a deprecated=message argument
  to the add_argument call for the deprecated argument. Beside logging the
  deprecation warning, the argument will behave as normal.
  """

  def __init__(self, *args, **kwargs):
    """Init override to extract the deprecated argument when it exists."""
    self.deprecated_message = kwargs.pop('deprecated', None)
    super(_DeprecatedAction, self).__init__(*args, **kwargs)

  def __call__(self, parser, namespace, values, option_string=None):
    """Log the message then defer to the parent action."""
    if self.deprecated_message:
      logging.warning('Argument %s is deprecated: %s', option_string,
                      self.deprecated_message)
    return super(_DeprecatedAction, self).__call__(
        parser, namespace, values, option_string=option_string)


def OptparseWrapCheck(desc, check_f, _option, opt, value):
  """Optparse adapter for type checking functionality."""
  try:
    return check_f(value)
  except ValueError:
    raise optparse.OptionValueError(
        'Invalid %s given: --%s=%s' % (desc, opt, value))


class Option(optparse.Option):
  """Subclass to implement path evaluation & other useful types."""

  _EXTRA_TYPES = ('path', 'gs_path')
  TYPES = optparse.Option.TYPES + _EXTRA_TYPES
  TYPE_CHECKER = optparse.Option.TYPE_CHECKER.copy()
  for t in _EXTRA_TYPES:
    TYPE_CHECKER[t] = functools.partial(OptparseWrapCheck, t, VALID_TYPES[t])


class FilteringOption(Option):
  """Subclass that supports Option filtering for FilteringOptionParser"""

  _EXTRA_ACTIONS = ('split_extend',)
  ACTIONS = Option.ACTIONS + _EXTRA_ACTIONS
  STORE_ACTIONS = Option.STORE_ACTIONS + _EXTRA_ACTIONS
  TYPED_ACTIONS = Option.TYPED_ACTIONS + _EXTRA_ACTIONS
  ALWAYS_TYPED_ACTIONS = (Option.ALWAYS_TYPED_ACTIONS + _EXTRA_ACTIONS)

  def take_action(self, action, dest, opt, value, values, parser):
    if action == 'split_extend':
      lvalue = value.split()
      values.ensure_value(dest, []).extend(lvalue)
    else:
      Option.take_action(self, action, dest, opt, value, values, parser)

    if value is None:
      value = []
    elif not self.nargs or self.nargs <= 1:
      value = [value]

    parser.AddParsedArg(self, opt, [str(v) for v in value])


class ColoredFormatter(logging.Formatter):
  """A logging formatter that can color the messages."""

  _COLOR_MAPPING = {
      'WARNING': terminal.Color.YELLOW,
      'ERROR': terminal.Color.RED,
  }

  def __init__(self, *args, **kwargs):
    """Initializes the formatter.

    Args:
      args: See logging.Formatter for specifics.
      kwargs: See logging.Formatter for specifics.
      enable_color: Whether to enable colored logging. Defaults
        to None, where terminal.Color will set to a sane default.
    """
    self.color = terminal.Color(enabled=kwargs.pop('enable_color', None))
    super(ColoredFormatter, self).__init__(*args, **kwargs)

  def format(self, record):
    """Formats |record| with color."""
    msg = super(ColoredFormatter, self).format(record)
    color = self._COLOR_MAPPING.get(record.levelname)
    return msg if not color else self.color.Color(color, msg)


class ChromiteStreamHandler(logging.StreamHandler):
  """A stream handler for logging."""


class BaseParser(object):
  """Base parser class that includes the logic to add logging controls."""

  DEFAULT_LOG_LEVELS = ('fatal', 'critical', 'error', 'warning', 'notice',
                        'info', 'debug')

  DEFAULT_LOG_LEVEL = 'info'
  ALLOW_LOGGING = True

  def __init__(self, **kwargs):
    """Initialize this parser instance.

    kwargs:
      logging: Defaults to ALLOW_LOGGING from the class; if given,
        add --log-level.
      default_log_level: If logging is enabled, override the default logging
        level.  Defaults to the class's DEFAULT_LOG_LEVEL value.
      log_levels: If logging is enabled, this overrides the enumeration of
        allowed logging levels.  If not given, defaults to the classes
        DEFAULT_LOG_LEVELS value.
      manual_debug: If logging is enabled and this is True, suppress addition
        of a --debug alias.  This option defaults to True unless 'debug' has
        been exempted from the allowed logging level targets.
      caching: If given, must be either a callable that discerns the cache
        location if it wasn't specified (the prototype must be akin to
        lambda parser, values:calculated_cache_dir_path; it may return None to
        indicate that it handles setting the value on its own later in the
        parsing including setting the env), or True; if True, the
        machinery defaults to invoking the class's FindCacheDir method
        (which can be overridden).  FindCacheDir $CROS_CACHEDIR, falling
        back to $REPO/.cache, finally falling back to $TMP.
        Note that the cache_dir is not created, just discerned where it
        should live.
        If False, or caching is not given, then no --cache-dir option will be
        added.
    """
    self.debug_enabled = False
    self.caching_group = None
    self.debug_group = None
    self.default_log_level = None
    self.log_levels = None
    self.logging_enabled = kwargs.get('logging', self.ALLOW_LOGGING)
    self.default_log_level = kwargs.get('default_log_level',
                                        self.DEFAULT_LOG_LEVEL)
    self.log_levels = tuple(x.lower() for x in
                            kwargs.get('log_levels', self.DEFAULT_LOG_LEVELS))
    self.debug_enabled = (not kwargs.get('manual_debug', False)
                          and 'debug' in self.log_levels)
    self.caching = kwargs.get('caching', False)
    self._cros_defaults = {}

  @staticmethod
  def PopUsedArgs(kwarg_dict):
    """Removes keys used by the base parser from the kwarg namespace."""
    parser_keys = ['logging', 'default_log_level', 'log_levels', 'manual_debug',
                   'caching']
    for key in parser_keys:
      kwarg_dict.pop(key, None)

  def SetupOptions(self):
    """Sets up standard chromite options."""
    # NB: All options here must go through add_common_argument_to_group.
    # You cannot use add_argument or such helpers directly.  This is to
    # support default values with subparsers.
    #
    # You should also explicitly add default=None here when you want the
    # default to be set up in the parsed option namespace.
    if self.logging_enabled:
      self.debug_group = self.add_argument_group('Debug options')
      self.add_common_argument_to_group(
          self.debug_group, '--log-level', choices=self.log_levels,
          default=self.default_log_level,
          help='Set logging level to report at.')
      self.add_common_argument_to_group(
          self.debug_group, '--log-format', action='store',
          default=constants.LOGGER_FMT,
          help='Set logging format to use.')
      # Backwards compat name.  We should delete this at some point.
      self.add_common_argument_to_group(
          self.debug_group, '--log_format', action='store',
          default=constants.LOGGER_FMT,
          help=argparse.SUPPRESS)
      self.add_common_argument_to_group(
          self.debug_group,
          '-v',
          '--verbose',
          action='store_const',
          const='info',
          dest='log_level',
          help='Alias for `--log-level=info`.')
      if self.debug_enabled:
        self.add_common_argument_to_group(
            self.debug_group, '--debug', action='store_const', const='debug',
            dest='log_level', help='Alias for `--log-level=debug`. '
            'Useful for debugging bugs/failures.')
      self.add_common_argument_to_group(
          self.debug_group, '--nocolor', action='store_false', dest='color',
          default=None,
          help='Do not use colorized output (or `export NOCOLOR=true`)')

    if self.caching:
      self.caching_group = self.add_argument_group('Caching Options')
      self.add_common_argument_to_group(
          self.caching_group, '--cache-dir', default=None, type='path',
          help='Override the calculated chromeos cache directory; '
          "typically defaults to '$REPO/.cache' .")

  def SetupLogging(self, opts):
    """Sets up logging based on |opts|."""
    value = opts.log_level.upper()
    logger = logging.getLogger()
    logger.setLevel(getattr(logging, value))
    formatter = ColoredFormatter(fmt=opts.log_format,
                                 datefmt=constants.LOGGER_DATE_FMT,
                                 enable_color=opts.color)

    # Only set colored formatter for ChromiteStreamHandler instances,
    # which could have been added by ScriptWrapperMain() below.
    chromite_handlers = [x for x in logger.handlers if
                         isinstance(x, ChromiteStreamHandler)]
    for handler in chromite_handlers:
      handler.setFormatter(formatter)

    logging.captureWarnings(True)

    return value

  def DoPostParseSetup(self, opts, args):
    """Method called to handle post opts/args setup.

    This can be anything from logging setup to positional arg count validation.

    Args:
      opts: optparse.Values or argparse.Namespace instance
      args: position arguments unconsumed from parsing.

    Returns:
      (opts, args), w/ whatever modification done.
    """
    for dest, default in self._cros_defaults.items():
      if not hasattr(opts, dest):
        setattr(opts, dest, default)

    if self.logging_enabled:
      value = self.SetupLogging(opts)
      if self.debug_enabled:
        opts.debug = (value == 'DEBUG')
      opts.verbose = value in ('INFO', 'DEBUG')

    if self.caching:
      path = os.environ.get(constants.SHARED_CACHE_ENVVAR)
      if path is not None and opts.cache_dir is None:
        opts.cache_dir = os.path.abspath(path)

      opts.cache_dir_specified = opts.cache_dir is not None
      if not opts.cache_dir_specified:
        func = self.FindCacheDir if not callable(self.caching) else self.caching
        opts.cache_dir = func(self, opts)
      if opts.cache_dir is not None:
        self.ConfigureCacheDir(opts.cache_dir)

    return opts, args

  @staticmethod
  def ConfigureCacheDir(cache_dir):
    if cache_dir is None:
      os.environ.pop(constants.SHARED_CACHE_ENVVAR, None)
      logging.debug('Removed cache_dir setting')
    else:
      os.environ[constants.SHARED_CACHE_ENVVAR] = cache_dir
      logging.debug('Configured cache_dir to %r', cache_dir)

  @classmethod
  def FindCacheDir(cls, _parser, _opts):
    logging.debug('Cache dir lookup.')
    return path_util.FindCacheDir()


@six.add_metaclass(attrs_freezer.Class)
class ArgumentNamespace(argparse.Namespace):
  """Class to mimic argparse.Namespace with value freezing support."""
  _FROZEN_ERR_MSG = 'Option values are frozen, cannot alter %s.'


# Note that because optparse.Values is not a new-style class this class
# must use the mixin rather than the metaclass.
class OptionValues(attrs_freezer.Mixin, optparse.Values):
  """Class to mimic optparse.Values with value freezing support."""
  _FROZEN_ERR_MSG = 'Option values are frozen, cannot alter %s.'

  def __init__(self, defaults, *args, **kwargs):
    attrs_freezer.Mixin.__init__(self)
    optparse.Values.__init__(self, defaults, *args, **kwargs)

    # Used by FilteringParser.
    self.parsed_args = None


PassedOption = collections.namedtuple(
    'PassedOption', ['opt_inst', 'opt_str', 'value_str'])


class FilteringParser(optparse.OptionParser, BaseParser):
  """Custom option parser for filtering options.

  Aside from adding a couple of types (path for absolute paths,
  gs_path for google storage urls, and log_level for logging level control),
  this additionally exposes logging control by default; if undesired,
  either derive from this class setting ALLOW_LOGGING to False, or
  pass in logging=False to the constructor.
  """

  DEFAULT_OPTION_CLASS = FilteringOption

  def __init__(self, usage=None, **kwargs):
    BaseParser.__init__(self, **kwargs)
    self.PopUsedArgs(kwargs)
    kwargs.setdefault('option_class', self.DEFAULT_OPTION_CLASS)
    optparse.OptionParser.__init__(self, usage=usage, **kwargs)
    self.SetupOptions()

  def add_common_argument_to_group(self, group, *args, **kwargs):
    """Adds the given option defined by args and kwargs to group."""
    return group.add_option(*args, **kwargs)

  def add_argument_group(self, *args, **kwargs):
    """Return an option group rather than an argument group."""
    return self.add_option_group(*args, **kwargs)

  def parse_args(self, args=None, values=None):
    # If no Values object is specified then use our custom OptionValues.
    if values is None:
      values = OptionValues(defaults=self.defaults)

    values.parsed_args = []

    opts, remaining = optparse.OptionParser.parse_args(
        self, args=args, values=values)
    return self.DoPostParseSetup(opts, remaining)

  def AddParsedArg(self, opt_inst, opt_str, value_str):
    """Add a parsed argument with attributes.

    Args:
      opt_inst: An instance of a raw optparse.Option object that represents the
                option.
      opt_str: The option string.
      value_str: A list of string-ified values dentified by OptParse.
    """
    self.values.parsed_args.append(PassedOption(opt_inst, opt_str, value_str))

  @staticmethod
  def FilterArgs(parsed_args, filter_fn):
    """Filter the argument by passing it through a function.

    Args:
      parsed_args: The list of parsed argument namedtuples to filter.  Tuples
        are of the form (opt_inst, opt_str, value_str).
      filter_fn: A function with signature f(PassedOption), and returns True if
        the argument is to be passed through.  False if not.

    Returns:
      A tuple containing two lists - one of accepted arguments and one of
      removed arguments.
    """
    removed = []
    accepted = []
    for arg in parsed_args:
      target = accepted if filter_fn(arg) else removed
      target.append(arg.opt_str)
      target.extend(arg.value_str)

    return accepted, removed


class ArgumentParser(BaseParser, argparse.ArgumentParser):
  """Custom argument parser for use by chromite.

  This class additionally exposes logging control by default; if undesired,
  either derive from this class setting ALLOW_LOGGING to False, or
  pass in logging=False to the constructor.
  """

  def __init__(self, usage=None, **kwargs):
    kwargs.setdefault('formatter_class', argparse.RawDescriptionHelpFormatter)
    BaseParser.__init__(self, **kwargs)
    self.PopUsedArgs(kwargs)
    argparse.ArgumentParser.__init__(self, usage=usage, **kwargs)
    self._SetupTypes()
    self.SetupOptions()
    self._RegisterActions()

  def _SetupTypes(self):
    """Register types with ArgumentParser."""
    for t, check_f in VALID_TYPES.items():
      self.register('type', t, check_f)
    for a, class_a in VALID_ACTIONS.items():
      self.register('action', a, class_a)

  def _RegisterActions(self):
    """Update the container's actions.

    This method builds out a new action class to register for each action type.
    The new action class allows handling the deprecated argument without any
    other changes to the argument parser logic. See _DeprecatedAction.
    """
    for action in _DEPRECATE_ACTIONS:
      current_class = self._registry_get('action', action, object)
      # Base classes for the new class. The _DeprecatedAction must be first to
      # ensure its method overrides are called first.
      bases = (_DeprecatedAction, current_class)
      try:
        self.register('action', action, type('deprecated-wrapper', bases, {}))
      except TypeError:
        # Method resolution order error. This occurs when the _DeprecatedAction
        # class is inherited multiple times, so we've already registered the
        # replacement class. The underlying _ActionsContainer gets passed
        # around, so this may get triggered in non-obvious ways.
        continue

  def add_common_argument_to_group(self, group, *args, **kwargs):
    """Adds the given argument to the group.

    This argument is expected to show up across the base parser and subparsers
    that might be added later on.  The default argparse module does not handle
    this scenario well -- it processes the base parser first (defaults and the
    user arguments), then it processes the subparser (defaults and arguments).
    That means defaults in the subparser will clobber user arguments passed in
    to the base parser!
    """
    default = kwargs.pop('default', None)
    kwargs['default'] = argparse.SUPPRESS
    action = group.add_argument(*args, **kwargs)
    self._cros_defaults.setdefault(action.dest, default)
    return action

  def parse_args(self, args=None, namespace=None):
    """Translates OptionParser call to equivalent ArgumentParser call."""
    # If no Namespace object is specified then use our custom ArgumentNamespace.
    if namespace is None:
      namespace = ArgumentNamespace()

    # Unlike OptionParser, ArgParser works only with a single namespace and no
    # args. Re-use BaseParser DoPostParseSetup but only take the namespace.
    namespace = argparse.ArgumentParser.parse_args(
        self, args=args, namespace=namespace)
    return self.DoPostParseSetup(namespace, None)[0]


class _ShutDownException(SystemExit):
  """Exception raised when user hits CTRL+C."""

  def __init__(self, sig_num, message):
    self.signal = sig_num
    # Setup a usage message primarily for any code that may intercept it
    # while this exception is crashing back up the stack to us.
    SystemExit.__init__(self, 128 + sig_num)
    self.args = (sig_num, message)

  def __str__(self):
    """Stringify this exception."""
    return self.args[1]


def _DefaultHandler(signum, _frame):
  # Don't double process sigterms; just trigger shutdown from the first
  # exception.
  signal.signal(signum, signal.SIG_IGN)
  raise _ShutDownException(
      signum, 'Received signal %i; shutting down' % (signum,))


def _RestartInChroot(cmd, chroot_args, extra_env):
  """Rerun inside the chroot.

  Args:
    cmd: Command line to run inside the chroot as a list of strings.
    chroot_args: Arguments to pass directly to cros_sdk (or None).
    extra_env: Dictionary of environmental variables to set inside the
        chroot (or None).
  """
  return cros_build_lib.run(cmd, check=False, enter_chroot=True,
                            chroot_args=chroot_args, extra_env=extra_env,
                            cwd=constants.SOURCE_ROOT).returncode


def RunInsideChroot(command=None, chroot_args=None):
  """Restart the current command inside the chroot.

  This method is only valid for any code that is run via ScriptWrapperMain.
  It allows proper cleanup of the local context by raising an exception handled
  in ScriptWrapperMain.

  Args:
    command: An instance of CliCommand to be restarted inside the chroot.
             |command| can be None if you do not wish to modify the log_level.
    chroot_args: List of command-line arguments to pass to cros_sdk, if invoked.
  """
  if cros_build_lib.IsInsideChroot():
    return

  # Produce the command line to execute inside the chroot.
  argv = sys.argv[:]
  argv[0] = path_util.ToChrootPath(argv[0])

  # Set log-level of cros_sdk to be same as log-level of command entering the
  # chroot.
  if chroot_args is None:
    chroot_args = []
  if command is not None:
    chroot_args += ['--log-level', command.options.log_level]

  raise ChrootRequiredError(argv, chroot_args)


def ReExec():
  """Restart the current command.

  This method is only valid for any code that is run via ScriptWrapperMain.
  It allows proper cleanup of the local context by raising an exception handled
  in ScriptWrapperMain.
  """
  # The command to exec.
  raise ExecRequiredError(sys.argv[:])


def ScriptWrapperMain(find_target_func, argv=None,
                      log_level=logging.DEBUG,
                      log_format=constants.LOGGER_FMT):
  """Function usable for chromite.script.* style wrapping.

  Note that this function invokes sys.exit on the way out by default.

  Args:
    find_target_func: a function, which, when given the absolute
      pathway the script was invoked via (for example,
      /home/ferringb/cros/trunk/chromite/bin/cros_sdk; note that any
      trailing .py from the path name will be removed),
      will return the main function to invoke (that functor will take
      a single arg- a list of arguments, and shall return either None
      or an integer, to indicate the exit code).
    argv: sys.argv, or an equivalent tuple for testing.  If nothing is
      given, sys.argv is defaulted to.
    log_level: Default logging level to start at.
    log_format: Default logging format to use.
  """
  if argv is None:
    argv = sys.argv[:]
  target = os.path.abspath(argv[0])
  name = os.path.basename(target)
  if target.endswith('.py'):
    target = os.path.splitext(target)[0]
  target = find_target_func(target)
  if target is None:
    print('Internal error detected- no main functor found in module %r.' %
          (name,), file=sys.stderr)
    sys.exit(100)

  # Set up basic logging information for all modules that use logging.
  # Note a script target may setup default logging in its module namespace
  # which will take precedence over this.
  logger = logging.getLogger()
  logger.setLevel(log_level)
  logger_handler = ChromiteStreamHandler()
  logger_handler.setFormatter(
      logging.Formatter(fmt=log_format, datefmt=constants.LOGGER_DATE_FMT))
  logger.addHandler(logger_handler)
  logging.captureWarnings(True)

  signal.signal(signal.SIGTERM, _DefaultHandler)

  ret = 1
  try:
    ret = target(argv[1:])
  except _ShutDownException as e:
    sys.stdout.flush()
    print('%s: Signaled to shutdown: caught %i signal.' % (name, e.signal),
          file=sys.stderr)
    sys.stderr.flush()
  except SystemExit as e:
    # Right now, let this crash through- longer term, we'll update the scripts
    # in question to not use sys.exit, and make this into a flagged error.
    raise
  except ChrootRequiredError as e:
    ret = _RestartInChroot(e.cmd, e.chroot_args, e.extra_env)
  except ExecRequiredError as e:
    logging.shutdown()
    # This does not return.
    os.execv(e.cmd[0], e.cmd)
  except Exception as e:
    sys.stdout.flush()
    print('%s: Unhandled exception:' % (name,), file=sys.stderr)
    sys.stderr.flush()
    raise
  finally:
    logging.shutdown()

  if ret is None:
    ret = 0
  sys.exit(ret)