Permalink
Cannot retrieve contributors at this time
Fetching contributors…
Cannot retrieve contributors at this time
| #!/usr/bin/env python | |
| import re | |
| import os | |
| import sys | |
| import yaml | |
| import subprocess | |
| import itertools | |
| from optparse import OptionParser | |
| import jujuclient | |
| import jujuclient.juju1.environment | |
| import jujuclient.juju2.environment | |
| JUJU_HOME = os.environ.get('JUJU_HOME', '') | |
| CONFIG_FILE = os.path.join(JUJU_HOME, 'dhx-rc.yaml') | |
| ALT_CONFIG_FILE = os.path.join(JUJU_HOME, 'debug-hooks-rc.yaml') | |
| DEFAULT_CONFIG = {'import_ids': [], | |
| 'uploads': {}, | |
| 'init': os.path.join(JUJU_HOME, 'debug-hooks-rc.d'), | |
| 'use_tmux_bindings': False, | |
| 'auto_sync': False, | |
| 'auto_restart': False, | |
| 'sync_exclude': ['*.pyc', '.*']} | |
| def call(*args): | |
| try: | |
| return subprocess.check_output(('/usr/bin/env',) + args, | |
| stderr=subprocess.STDOUT) | |
| except subprocess.CalledProcessError as e: | |
| sys.stderr.write(e.output) | |
| sys.exit(e.returncode) | |
| JUJU_ENV = call('juju', 'switch').rstrip('\n') | |
| JUJU_VERSION = call('juju', 'version') | |
| JUJU_2 = JUJU_VERSION.startswith('2.') | |
| class CaseInsensitiveDict(dict): | |
| def __init__(self, *args, **kwargs): | |
| super(CaseInsensitiveDict, self).__init__(*args, **kwargs) | |
| for key in self.keys(): | |
| if isinstance(self[key], dict): | |
| self[key] = CaseInsensitiveDict(self[key]) | |
| def __getitem__(self, key): | |
| if key.lower() in self: | |
| key = key.lower() | |
| return super(CaseInsensitiveDict, self).__getitem__(key) | |
| def get(self, key, default=None): | |
| if key.lower() in self: | |
| key = key.lower() | |
| return super(CaseInsensitiveDict, self).get(key, default) | |
| def show_help(option, opt, value, parser): | |
| print 'usage: juju dhx [options] [<unit-or-service>]' | |
| print 'purpose: Enhanced interactive remote hook debugging session, using tmux' | |
| print 'options:' | |
| print '-c, --config (= "$JUJU_HOME/dhx-rc.yaml")' | |
| print ' yaml file to use for customizing the debug-hooks environment' | |
| print '-e, --environment (= current)' | |
| print ' juju environment to operate in' | |
| print '-i, --import-id AUTH_ID' | |
| print ' import an auth key for a shared debug session from Launchpad (you' | |
| print ' can also add IDs to the config file if you have some that you always' | |
| print ' want to import)' | |
| print '-j, --join' | |
| print ' join an existing remove debug session (useful for paired debugging)' | |
| print ' does not use the Juju identity, so one of your ssh identities' | |
| print ' must have been imported (e.g., with import_ids, below)' | |
| print ' note that your customizations will not be made to joined sessions' | |
| print '-r, --retry, --restart' | |
| print ' restart the failed hook, such that the session can debug it' | |
| print '-R, --no-retry, --no-restart' | |
| print ' do not restart the failed hook, even if auto_restart is true or' | |
| print ' --restart is given' | |
| print '-s, --sync' | |
| print ' sync changes to the charm from the remote unit back locally' | |
| print ' this ensures that changes made while debugging the charm' | |
| print ' are preserved' | |
| print '-S, --no-sync' | |
| print ' do not sync changes, even if auto_sync is true or --sync is given' | |
| print 'config file options:' | |
| print 'use_tmux_bindings' | |
| print ' create an empty .tmux.conf to reset tmux to its default key bindings' | |
| print ' instead of using the screen bindings (note: this option has no effect' | |
| print ' if a custom .tmux.conf file is created via the uploads or init options)' | |
| print 'init' | |
| print ' file or directory name containing script(s) to be run once to setup' | |
| print ' a new machine; will create or append to a .bashrc file on the machine' | |
| print 'uploads' | |
| print ' a list of (or map of local to remote) files or directories to upload' | |
| print ' to the remote machine (this is only done once, and requires rsync)' | |
| print 'import_ids' | |
| print ' list of Launchpad user IDs to import into Juju to give access to' | |
| print ' the remote debugging session (they should then use --join)' | |
| print 'auto_restart' | |
| print ' always assume the --restart option' | |
| print 'auto_sync' | |
| print ' always assume the --sync option' | |
| print 'sync_exclude' | |
| print ' list of file patterns for --sync to exclude (see rysnc --exclude)' | |
| print 'Note that the init script(s) will be run once, the first time the machine' | |
| print 'is connected to. If you wish to have setup that runs each time, you will' | |
| print 'need to replace or append to the .bashrc file. Also note that the init' | |
| print 'script(s) will block the login process, so if you have setup tasks that' | |
| print 'take a while to run, such as installing packages, you may want to run' | |
| print 'them in the background.' | |
| sys.exit() | |
| def show_desc(option, opt, value, parser): | |
| if os.path.basename(__file__) == 'juju-debug-hooks-ext': | |
| print 'Alias for juju-dhx' | |
| else: | |
| print 'Enhanced interactive remote hook debugging session, using tmux' | |
| sys.exit(0) | |
| def parse_args(): | |
| parser = OptionParser(add_help_option=False) | |
| parser.add_option('-c', '--config', dest='config_file', default=None) | |
| parser.add_option('-e', '--environment', dest='env', default=JUJU_ENV) | |
| parser.add_option('-i', '--import-id', dest='import_ids', action='append', default=[]) | |
| parser.add_option('-j', '--join', dest='join', action='store_true') | |
| parser.add_option('-r', '--retry', '--restart', dest='restart', action='store_true') | |
| parser.add_option('-R', '--no-retry', '--no-restart', dest='no_restart', action='store_true') | |
| parser.add_option('-s', '--sync', dest='sync', action='store_true') | |
| parser.add_option('-S', '--no-sync', dest='no_sync', action='store_true') | |
| parser.add_option('-h', '--help', action='callback', callback=show_help) | |
| parser.add_option('-d', '--description', action='callback', callback=show_desc) | |
| opts = parser.parse_args() | |
| parser.values.config_file = find_config_file(parser.values.config_file) | |
| return opts | |
| def find_config_file(config_file): | |
| if config_file is not None: | |
| config_file = os.path.expanduser(config_file) | |
| if os.path.exists(config_file): | |
| return config_file | |
| else: | |
| if os.path.exists(CONFIG_FILE): | |
| return CONFIG_FILE | |
| elif os.path.exists(ALT_CONFIG_FILE): | |
| return ALT_CONFIG_FILE | |
| return None | |
| def customize(unit_name, unit, opts): | |
| config = dict(DEFAULT_CONFIG) | |
| if opts.config_file: | |
| with open(opts.config_file) as fp: | |
| config = yaml.load(fp) | |
| import_ids = config.get('import_ids', []) + opts.import_ids | |
| do_import_ids(unit, import_ids, opts) | |
| if check_customized(unit_name, unit, opts): | |
| return config | |
| sys.stdout.write('Customizing machine...') | |
| sys.stdout.flush() | |
| uploads = config.get('uploads', {}) | |
| if isinstance(uploads, (list, tuple)): | |
| uploads = {u: '' for u in uploads} | |
| for local, remote in uploads.iteritems(): | |
| upload(unit_name, local, remote) | |
| if os.path.exists(os.path.expanduser(config.get('init', ''))): | |
| upload(unit_name, config['init'], '') | |
| upload(unit_name, os.path.join(os.path.dirname(__file__), 'dhx-init.sh'), | |
| '.dhx.init') | |
| call('juju', 'ssh', unit_name, | |
| 'echo "DHX_USER_INIT=\'%s\' ./.dhx.init" >> .bashrc' | |
| % os.path.basename(config.get('init', ''))) | |
| if config.get('use_tmux_bindings', False): | |
| call('juju', 'ssh', unit_name, 'sudo touch .tmux.conf') | |
| print 'done.' | |
| set_customized(unit_name, unit, opts) | |
| return config | |
| def check_customized(unit_name, unit, opts): | |
| env = get_env(opts) | |
| machine = unit.get('Machine') | |
| annotations = env.get_annotation(machine, 'machine')['Annotations'] or {} | |
| return 'dhx-customized' in annotations | |
| def set_customized(unit_name, unit, opts): | |
| env = get_env(opts) | |
| machine = unit.get('Machine') | |
| env.set_annotation(machine, 'machine', {'dhx-customized': 'true'}) | |
| def do_import_ids(unit, import_ids, opts): | |
| if not import_ids: | |
| return | |
| env = get_env(opts) | |
| machine = unit.get('Machine') | |
| annotations = env.get_annotation(machine, 'machine')['Annotations'] or {} | |
| import_ids = ['lp:{}'.format(iid) if ':' not in iid else iid | |
| for iid in import_ids] | |
| if annotations.get('import-ids') != ','.join(import_ids): | |
| cmd = ['juju'] | |
| if JUJU_2: | |
| cmd.append('import-ssh-key') | |
| else: | |
| cmd.extend(['authorized-keys', 'import']) | |
| cmd.extend(import_ids) | |
| call(*cmd) | |
| env.set_annotation(machine, 'machine', {'import-ids': ','.join(import_ids)}) | |
| def get_env(opts): | |
| if hasattr(get_env, '_env'): | |
| return get_env._env | |
| if JUJU_2: | |
| get_env._env = jujuclient.juju2.environment.Environment.connect(opts.env) | |
| else: | |
| get_env._env = jujuclient.juju1.environment.Environment.connect(opts.env) | |
| return get_env._env | |
| def _unit_errored(unit): | |
| if JUJU_2: | |
| return unit['workload-status']['status'] == 'error' | |
| else: | |
| return unit['AgentState'] == 'error' | |
| def find_errored_unit(units): | |
| if not units: | |
| return None | |
| for name, unit in units.items(): | |
| if _unit_errored(unit): | |
| return name | |
| return None | |
| def units_for_service(services, service_name): | |
| units = {} | |
| service = services[service_name] | |
| if service.get('Units'): | |
| units = service['Units'] | |
| if 'SubordinateTo' in service: | |
| for sub_to in service['SubordinateTo']: | |
| units.update(services[sub_to].get('Units', {})) | |
| units = {n: u for n, u in units.items() | |
| if u.get('PublicAddress')} | |
| return units | |
| def fail(msg): | |
| sys.stderr.write('%s\n' % msg) | |
| sys.exit(1) | |
| def current_charm_name(): | |
| if not os.path.exists('metadata.yaml'): | |
| return None | |
| with open('metadata.yaml') as fp: | |
| metadata = yaml.safe_load(fp) | |
| return metadata['name'] | |
| def parse_charm_id(charm_id): | |
| charm_match = re.match(r'local:[^/]*/(.*)-\d+', charm_id) | |
| if charm_match: | |
| return charm_match.group(1) | |
| return None | |
| def get_current_service_units(services): | |
| charm_name = current_charm_name() | |
| if charm_name: | |
| for service_name, service in services.items(): | |
| if parse_charm_id(service['Charm']) == charm_name: | |
| return units_for_service(services, service_name) | |
| return None | |
| def choose_unit_from(all_units, preferred_units=None): | |
| if len(all_units) == 1: | |
| return all_units.items()[0] | |
| default_unit = find_errored_unit(preferred_units) | |
| if not default_unit: | |
| default_unit = find_errored_unit(all_units) | |
| if not default_unit and preferred_units: | |
| default_unit = sorted(preferred_units.keys())[0] | |
| if not default_unit: | |
| default_unit = sorted(all_units.keys())[0] | |
| unit_names = sorted(all_units.keys()) | |
| print 'Units:' | |
| for i, unit_name in enumerate(unit_names): | |
| unit_err = '' | |
| if _unit_errored(all_units[unit_name]): | |
| unit_err = ' (error)' | |
| print ' %d: %s%s' % (i, unit_name, unit_err) | |
| selection = raw_input('Select a unit by number or name: %s' % | |
| ('[%s] ' % default_unit or '')) | |
| if selection.isdigit(): | |
| unit = unit_names[int(selection)] | |
| elif selection == '': | |
| unit = default_unit | |
| else: | |
| unit = selection | |
| if unit not in all_units: | |
| fail('Invalid unit: %s' % unit) | |
| return unit, all_units[unit] | |
| def get_all_units(status): | |
| services = CaseInsensitiveDict(status.get('applications', status.get('Services'))) | |
| units = {n: CaseInsensitiveDict(u, Charm=s.get('Charm')) | |
| for s in services.values() if s.get('Units') | |
| for n, u in s['Units'].iteritems()} | |
| subordinates = {} | |
| for name, unit in units.iteritems(): | |
| subs = unit.get('Subordinates') or {} | |
| for sub_name, sub_unit in subs.iteritems(): | |
| sub_unit['Machine'] = unit['Machine'] | |
| subordinates[sub_name] = sub_unit | |
| units.update(subordinates) | |
| return units | |
| def choose_unit(args, opts): | |
| unit = args.pop(0) if args else None | |
| status = CaseInsensitiveDict(get_env(opts).status()) | |
| services = status.get('applications', status.get('Services')) | |
| units = get_all_units(status) | |
| if not units: | |
| fail('No units available') | |
| if unit in units: # explicit unit given | |
| return unit, units[unit] | |
| elif unit in services: # explicit service given | |
| units = units_for_service(services, service_name=unit) | |
| if not units: | |
| fail('No units available for service: %s' % unit) | |
| return choose_unit_from(units) | |
| elif unit: | |
| fail('Invalid unit: %s' % unit) | |
| else: | |
| preferred_units = get_current_service_units(services) | |
| return choose_unit_from(units, preferred_units) | |
| def upload(unit, local, remote): | |
| call('rsync', '-Wa', '-e', 'juju ssh --pty=false %s' % unit, | |
| os.path.expanduser(local), ':%s' % remote) | |
| if __name__ == '__main__': | |
| opts, args = parse_args() | |
| if opts.join: | |
| if len(args) == 0: | |
| fail('You must provide an address') | |
| addr = args[0] | |
| os.execvp('ssh', ['ssh', '-t', 'ubuntu@%s' % addr, 'sudo tmux attach']) | |
| else: | |
| unit_name, unit = choose_unit(args, opts) | |
| config = customize(unit_name, unit, opts) | |
| model_opt = '-m' if JUJU_2 else '-e' | |
| if (opts.restart or config.get('auto_restart', False)) and not opts.no_restart: | |
| subprocess.Popen( | |
| 'sleep 2 ; juju resolved --retry %s %s "%s"' % (model_opt, opts.env, unit_name), | |
| stdout=subprocess.PIPE, stderr=subprocess.STDOUT, shell=True) | |
| result = subprocess.call( | |
| ['/usr/bin/env', 'juju', 'debug-hooks', model_opt, opts.env, unit_name] + args) | |
| if result != 0: | |
| sys.exit(result) | |
| if (opts.sync or config.get('auto_sync', False)) and not opts.no_sync: | |
| local_charm = current_charm_name() | |
| remote_charm = parse_charm_id(unit['Charm']) | |
| if remote_charm == local_charm: | |
| print 'Syncing remote changes...' | |
| excludes = list(itertools.chain.from_iterable( | |
| [('--exclude', e) for e in config.get('sync_exclude', [])])) | |
| os.execvp('juju', | |
| ['juju', 'sync-charm', '-y', unit_name, '.'] + excludes) | |
| elif opts.sync: | |
| print "Unable to sync remote changes: remote charm doesn't match current directory" |