Пример #1
0
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))
Пример #2
0
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)