py/instalog/cli.py

#!/usr/bin/env python3
#
# Copyright 2016 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.

from __future__ import print_function

import argparse
import json
import logging
import os
import signal
import sys
import tarfile

from cros.factory.instalog import core
from cros.factory.instalog import daemon_utils
from cros.factory.instalog import instalog_common
from cros.factory.instalog import log_utils
from cros.factory.instalog.utils import file_utils
from cros.factory.instalog.utils import sync_utils
from cros.factory.instalog.utils import type_utils

from cros.factory.instalog.external import jsonrpclib
from cros.factory.instalog.external import yaml


# The default number of seconds to wait before giving up on a flush.
_DEFAULT_FLUSH_TIMEOUT = 30
_DEFAULT_STOP_TIMEOUT = 10


class InstalogService(daemon_utils.Daemon):
  """Represents the Instalog daemon service."""

  def __init__(self, config, logging_level):
    self._config = config
    self._logging_level = logging_level
    self._core = None
    super(InstalogService, self).__init__(
        pidfile=config['instalog']['pid_file'])

  def _SignalHandler(self, signal_num, frame):
    """Signal handler to stop Instalog on SIGINT or SIGTERM."""
    del frame
    logging.debug('_SignalHandler called with signalnum=%s', signal_num)
    if signal_num not in [signal.SIGINT, signal.SIGTERM]:
      return
    if self._core:
      # No need for a lock since _SignalHandler will only ever be called from
      # Instalog's main thread.
      signal_string = 'SIGINT' if signal_num == signal.SIGINT else 'SIGTERM'
      logging.warning('%s detected, stopping', signal_string)
      self._core.Stop()
      self._core = None

  def _InitLogging(self, foreground):
    """Sets up logging."""
    handlers = []

    # Save logging calls to log file.
    log_file = self._config['instalog']['log_file']
    file_utils.TryMakeDirs(os.path.dirname(log_file))
    handlers.append(log_utils.GetFileHandler(log_file, self._logging_level))

    # Output logging calls to console when foreground is set.
    if foreground:
      handlers.append(log_utils.GetStreamHandler(self._logging_level))

    log_utils.InitLogging(handlers)

  def Run(self, foreground, rpc_ready=None):
    """Starts Instalog."""
    self._InitLogging(foreground)

    signal.signal(signal.SIGINT, self._SignalHandler)
    signal.signal(signal.SIGTERM, self._SignalHandler)

    self._core = core.Instalog(
        node_id=self._config['instalog']['node_id'],
        data_dir=self._config['instalog']['data_dir'],
        cli_hostname=self._config['instalog']['cli_hostname'],
        cli_port=self._config['instalog']['cli_port'],
        buffer_plugin=self._config['buffer'],
        input_plugins=self._config['input'],
        output_plugins=self._config['output'])
    # After the core initialized, the RPC server is ready.
    if rpc_ready:
      rpc_ready.set()
    self._core.Run()


class InstalogCLI:
  """Represents the CLI interface used to control Instalog."""

  def __init__(self, args):
    # Read config file.
    config_path = self._LocateConfigFile(args.config)
    if config_path is None:
      exit('No config file found')
    with open(config_path) as f:
      config = yaml.load(f)
    self._CheckDataDir(config)

    # logging.WARNING = 30, logging.INFO = 20, logging.DEBUG = 10
    logging_level = logging.INFO - ((args.verbose - args.quiet) * 10)

    self._service = InstalogService(config, logging_level)
    self._core = jsonrpclib.Server(
        'http://%s:%s' % (config['instalog']['cli_hostname'],
                          config['instalog']['cli_port']))

    if args.cmd == 'start':
      self.Start(args.foreground)
    elif args.cmd == 'stop':
      self.Stop(args.timeout)
    elif args.cmd == 'restart':
      self.Restart()
    elif args.cmd == 'status':
      self.Status()
    elif args.cmd == 'inspect':
      self.Inspect(args.plugin_id, args.json_path)
    elif args.cmd == 'flush':
      self.Flush(args.plugin_id, args.timeout)
    elif args.cmd == 'archive':
      self.Archive(config_path, config['instalog']['data_dir'],
                   args.archive_path, args.details)
    elif args.cmd == 'progress':
      self.Progress(args.plugin_id, args.details)

  def _LocateConfigFile(self, user_path):
    """Locates the config file that should be used by Instalog."""
    if user_path:
      return user_path
    paths = [
        os.path.join(os.getcwd(), 'instalog.yaml'),
        os.path.join(os.path.dirname(os.path.realpath(__file__)),
                     'instalog.yaml'),
        os.path.join(os.path.expanduser('~'), '.instalog.yaml'),
        os.path.join(os.sep, 'etc', 'instalog.yaml'),
        os.path.join(os.sep, 'run', 'instalog.yaml')]
    for path in paths:
      logging.debug('Checking %s for config file...', path)
      if os.path.exists(path):
        logging.info('Config file found at %s', path)
        return path
    return None

  def _CheckDataDir(self, config):
    data_dir = config['instalog']['data_dir']
    if not os.path.exists(data_dir):
      os.makedirs(data_dir)
    instalog_dir = instalog_common.INSTALOG_DIR
    for path, unused_dirs, unused_files in os.walk(
        instalog_dir, followlinks=True):
      if not os.path.islink(path) and os.path.samefile(path, data_dir):
        print('You should not put the data_dir in the Instalog source code')
        sys.exit(1)

  def Restart(self):
    """Restarts the daemon."""
    self.Stop(_DEFAULT_STOP_TIMEOUT)
    self.Start(False)

  def Start(self, foreground):
    """Starts the daemon.

    Args:
      foreground: Does not detach the daemon.
    """
    print('Starting...')
    if not self._service.Start(foreground):
      return
    if foreground:
      return

    # First, wait for the daemon process to start.
    try:
      sync_utils.WaitFor(self._service.IsRunning, 10)
    except type_utils.TimeoutError:
      print('Daemon could not be brought up, check the logs')
      sys.exit(1)

    def TryIsUp():
      try:
        # Perform the real check to see if Instalog is up internally.
        return self._core.IsUp()
      except Exception:
        raise type_utils.TimeoutError('Could not call core IsUp')

    try:
      print('Waiting for the core is up...')
      if sync_utils.WaitFor(TryIsUp, 10):
        print('DONE')
        return
    except type_utils.TimeoutError:
      pass
    print('Daemon could not be brought up, check the logs')
    sys.exit(1)

  def Stop(self, timeout):
    """Stops the daemon."""
    # First, send the "stop" instruction to the daemon.
    print('Stopping...')
    try:
      self._core.Stop()
    except Exception:
      print('Could not connect to daemon, is it running?')
      sys.exit(1)

    # Then, wait for the process to come down.
    try:
      sync_utils.WaitFor(self._service.IsStopped, timeout)
    except type_utils.TimeoutError:
      print('Still shutting down?')
      sys.exit(1)
    else:
      print('DONE')

  def Status(self):
    """Prints the status of the daemon."""
    running = self._service.IsRunning()
    if running:
      up = self._core.IsUp()
      print('UP' if up else 'STARTING')
    else:
      print('DOWN')

  def Inspect(self, plugin_id, json_path):
    """Inspects the store of a given plugin."""
    success, value = self._core.Inspect(plugin_id, json_path)
    print(value)
    if not success:
      sys.exit(1)

  def Flush(self, plugin_id, timeout):
    """Flushes the given plugin with given timeout."""
    success, value = self._core.Flush(plugin_id, timeout)
    print(json.dumps(value))
    if not success:
      sys.exit(1)

  def Archive(self, config_path, data_dir, archive_path, details):
    """Archives the whole Instalog."""
    if self._service.IsRunning():
      print('Is the Instalog running? You need to stop the Instalog first')
      sys.exit(1)
    if os.path.isdir(archive_path):
      archive_path = os.path.join(archive_path, 'archived_instalog.tar.gz')
    if not os.path.isdir(os.path.dirname(archive_path)):
      print('The directory of `%s` does not exist' %
            os.path.realpath(archive_path))
      sys.exit(1)

    print('Archiving to %s ...' % os.path.realpath(archive_path))
    with tarfile.open(archive_path, 'w') as tar:
      data_dir = os.path.realpath(data_dir)
      instalog_dir = instalog_common.INSTALOG_DIR
      instalog_parent_dir = instalog_common.INSTALOG_PARENT_DIR
      instalog_virtual_env_dir = instalog_common.INSTALOG_VIRTUAL_ENV_DIR

      if os.path.exists(data_dir):
        print('Archiving data_dir from %s' % os.path.realpath(data_dir))
        tar.add(data_dir, 'data')
        print('Archiving config file from %s' % os.path.realpath(config_path))
        tar.add(config_path, 'instalog.yaml')
      if details >= 1:
        def VirtualEnvFilter(tarinfo):
          if tarinfo.name == 'instalog/virtual_env':
            return None
          return tarinfo
        print('Archiving Instalog source code')
        tar.add(instalog_dir, 'instalog', filter=VirtualEnvFilter)
        tar.add(os.path.join(instalog_parent_dir, 'utils'), 'utils')
        tar.add(os.path.join(instalog_parent_dir, 'testlog'), 'testlog')
        tar.add(os.path.join(instalog_parent_dir, 'external'), 'external')
      if details >= 2:
        if os.path.exists(instalog_virtual_env_dir):
          print('Archiving virtual_env')
          tar.add(instalog_virtual_env_dir, 'instalog/virtual_env')
    print('DONE')

  def Progress(self, plugin_id, details):
    """Shows the progress of output plugins"""
    progress_dict = self._core.GetAllProgress(details)
    for name in sorted(progress_dict):
      if plugin_id is None or name.startswith(plugin_id):
        completed, total = progress_dict[name]
        print('%s completed %d of %d events, and remaining %d events' %
              (name, completed, total, total - completed))


def main():
  parser = argparse.ArgumentParser()
  parser.add_argument(
      '--config', '-c',
      help='config file path; by default, searches: \n'
           '$PWD/instalog.yaml py/instalog/instalog.yaml '
           '~/.instalog.yaml /etc/instalog.yaml /run/instalog.yaml')
  parser.add_argument(
      '--verbose', '-v', action='count', default=0,
      help='increase verbosity')
  parser.add_argument(
      '--quiet', '-q', action='count', default=0,
      help='decrease verbosity')

  subparsers = parser.add_subparsers(title='commands')

  start_parser = subparsers.add_parser('start', help='start Instalog')
  start_parser.set_defaults(cmd='start')
  start_parser.add_argument(
      '--no-daemon', '-n', dest='foreground', action='store_true',
      help='keep in foreground')

  stop_parser = subparsers.add_parser('stop', help='stop Instalog')
  stop_parser.set_defaults(cmd='stop')
  stop_parser.add_argument(
      '--timeout', '-w', type=float,
      required=False, default=_DEFAULT_STOP_TIMEOUT,
      help='time to wait before giving up')

  restart_parser = subparsers.add_parser('restart', help='restart Instalog')
  restart_parser.set_defaults(cmd='restart')

  status_parser = subparsers.add_parser('status', help='print Instalog status')
  status_parser.set_defaults(cmd='status')

  inspect_parser = subparsers.add_parser('inspect', help='inspect plugin store')
  inspect_parser.set_defaults(cmd='inspect')
  inspect_parser.add_argument(
      'plugin_id', type=str, help='ID of plugin to inspect')
  inspect_parser.add_argument(
      'json_path', type=str, nargs='?', default='.',
      help='path of store JSON to print')

  flush_parser = subparsers.add_parser('flush', help='flush plugin')
  flush_parser.set_defaults(cmd='flush')
  flush_parser.add_argument(
      '--timeout', '-w', type=float,
      required=False, default=_DEFAULT_FLUSH_TIMEOUT,
      help='time to wait before giving up')
  flush_parser.add_argument(
      'plugin_id', type=str, nargs='?', default=None,
      help='ID of plugin to flush')

  archive_parser = subparsers.add_parser('archive', help='archive the Instalog')
  archive_parser.set_defaults(cmd='archive')
  archive_parser.add_argument(
      '--output', '-o', dest='archive_path', type=str,
      required=False, default='.',
      help='path to put the archive file')
  archive_parser.add_argument(
      '--details', '-d', action='count', default=0,
      help='archive more details (instalog code / virtual_env)')

  progress_parser = subparsers.add_parser(
      'progress', help='print the progress of plugin')
  progress_parser.set_defaults(cmd='progress')
  progress_parser.add_argument(
      'plugin_id', type=str, nargs='?', default=None,
      help='ID of plugin\'s progress to print')
  progress_parser.add_argument(
      '--details', '-d', action='count', default=0,
      help='print more details')

  args = parser.parse_args()

  InstalogCLI(args)

if __name__ == '__main__':
  main()