|  | #!/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. | 
|  |  | 
|  | 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: | 
|  | sys.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() |