def main(progname=sys.argv[0]): args = _get_args() profile = LsiProfile.from_args(args) if args.host is not None: return get_host(args.host) # Either of these directives should force a refresh. latest = args.latest or args.refresh_only entries = get_entries(latest, profile.filters, profile.exclude) if args.version is True: _print_version() if args.refresh_only is True: print('Refreshed cache') return sort_by = args.sort_by or "name" if args.get is not None: _copy_from(entries, remote_path=args.get[0], local_path=args.get[1], profile=profile) elif args.put is not None: _copy_to(entries, local_path=args.put[0], remote_path=args.put[1], profile=profile) elif args.ssh is True or \ args.username is not None or args.identity_file is not None or \ args.profile is not None or profile.command is not None: _run_ssh( entries=entries, username=profile.username, idfile=profile.identity_file, no_prompt=profile.no_prompt, command=profile.command, show=args.show, only=args.only, sort_by=sort_by, limit=args.limit, tunnel=args.tunnel ) elif args.sep is not None: for e in entries: print(e.repr_as_line(additional_columns=args.show, only_show=args.only, sep=args.sep)) elif args.attributes is True: attribs = HostEntry.list_attributes() print('The following attributes are available: {}' .format(', '.join(attribs))) else: entries = HostEntry.sort_by(entries, sort_by) if args.limit is not None: entries = entries[:args.limit] print(HostEntry.render_entries(entries, additional_columns=args.show, only_show=args.only)) print('%s matching entries.' % len(entries))
def _run_ssh(entries, username, idfile, no_prompt=False, command=None, show=None, only=None, sort_by=None, limit=None, tunnel=None): """ Lets the user choose which instance to SSH into. :param entries: The list of host entries. :type entries: [:py:class:`HostEntry`] :param username: The SSH username to use. Defaults to current user. :type username: ``str`` or ``NoneType`` :param idfile: The identity file to use. Optional. :type idfile: ``str`` or ``NoneType`` :param no_prompt: Whether to disable confirmation for SSH command. :type no_prompt: ``bool`` :param command: SSH command to run on matching instances. :type command: ``str`` or ``NoneType`` :param show: Instance attributes to show in addition to defaults. :type show: ``NoneType`` or ``list`` of ``str`` :param only: If not ``None``, will *only* show these attributes. :type only: ``NoneType`` or ``list`` of ``str`` :param sort_by: What to sort columns by. By default, sort by 'name'. :type sort_by: ``str`` :param limit: At most how many results to show. :type limit: ``int`` or ``NoneType`` """ _print_entries = True _print_help = False if len(entries) == 0: exit('No entries matched the filters.') if no_prompt is True and command is not None: return _run_ssh_command(entries, username, idfile, command, tunnel) elif len(entries) == 1: if command is None: return _connect_ssh(entries[0], username, idfile, tunnel) else: return _run_ssh_command(entries, username, idfile, command, tunnel) elif command is not None: print(HostEntry.render_entries(entries, additional_columns=show, only_show=only, numbers=True)) if no_prompt is False: get_input("Press enter to run command {} on the {} " "above machines (Ctrl-C to cancel)" .format(cyan(command), len(entries))) return _run_ssh_command(entries, username, idfile, command, tunnel) else: while True: if sort_by is not None: entries = HostEntry.sort_by(entries, sort_by) if limit is not None: entries = entries[:limit] if _print_entries is True: print(HostEntry.render_entries(entries, additional_columns=show, only_show=only, numbers=True)) print('%s matching entries.' % len(entries)) _print_entries = False if _print_help is True: cmd_str = green(command) if command is not None else 'none set' msg = COMMANDS_STRING.format(username=username or 'none set', idfile=idfile or 'none set', cur_cmd=cmd_str) print(msg) _print_help = False elif command is not None: print('Set to run ssh command: %s' % cyan(command)) msg = 'Enter command (%s for help, %s to quit): ' % (cyan('h'), cyan('q')) choice = get_input(msg) if isinstance(choice, int): if 0 <= choice <= len(entries): break else: msg = 'Invalid number: must be between 0 and %s' print(msg % (len(entries) - 1)) elif choice == 'x': if command is None: print('No command has been set. Set command with `c`') else: return _run_ssh_command(entries, username, idfile, command, tunnel) elif choice == 'h': _print_help = True elif choice in ['q', 'quit', 'exit']: print('bye!') return else: # All of these commands take one or more arguments, so the # split length must be at least 2. commands = choice.split() if len(commands) < 2: print(yellow('Unknown command "%s".' % choice)) else: cmd = commands[0] if cmd in ['u', 'i', 'p']: if cmd == 'u': username = commands[1] elif cmd == 'i': _idfile = commands[1] if not os.path.exists(_idfile): print(yellow('No such file: %s' % _idfile)) continue idfile = _idfile elif cmd == 'p': p = commands[1] try: profile = LsiProfile.load(p) _username = profile.username _idfile = expanduser(profile.identity_file) except LsiProfile.LoadError: print(yellow('No such profile: %s' % repr(p))) continue username = _username idfile = _idfile print('username: %s' % green(repr(username))) print('identity file: %s' % green(repr(idfile))) elif cmd == 'f': entries = filter_entries(entries, commands[1:], []) _print_entries = True elif cmd == 'e': entries = filter_entries(entries, [], commands[1:]) _print_entries = True elif cmd == 'c': command = ' '.join(commands[1:]) elif cmd == 'limit': try: limit = int(commands[1]) _print_entries = True except ValueError: print(yellow('Invalid limit (must be an integer)')) elif cmd == 'sort': sort_by = commands[1] if sort_by not in show: show.append(sort_by) _print_entries = True elif cmd == 'show': if show is None: show = commands[1:] else: show.extend(commands[1:]) _print_entries = True else: print(yellow('Unknown command "%s".' % cmd)) return _connect_ssh(entries[choice], username, idfile, tunnel)