SSH Chooser: A simple script to extract hosts from configuration file and produce a navigation page

ssh_chooser.py

#!/usr/bin/env python3

# Install package windows-curses (via PyPI) before use on Windows platforms
import curses

import os
import re
import sys


config = os.path.expanduser('~/.ssh/config')


def print_center_core(scr, lines):
    baseline = (curses.LINES - len(lines) - 1) // 2
    for i in range(len(lines)):
        scr.addstr(baseline + i, (curses.COLS - len(lines[i]) - 1) // 2, lines[i])
    scr.move(0, 0)


def print_center(scr, text):
    print_center_core(scr, text.split('\n'))


def to_block(text):
    lines = text.split('\n')
    max_length = max([len(line) for line in lines])
    return ['{line:<{max_length}}'.format(line=line, max_length=max_length) for line in lines]


def get_hosts():
    res = []
    try:
        with open(config, encoding='utf-8') as file:
            content = file.read()
            exp = re.compile('^\\s*Host\\s+([^\\*]*?)$', flags=re.MULTILINE)
            for match in exp.finditer(content):
                res.append(match.expand('\\1'))
    except FileNotFoundError:
        pass
    return res


def error_no_host(scr):
    scr.clear()
    print_center(scr, 'ERROR\n\nNo configured host available')
    scr.refresh()
    scr.getch()


def choose_host(scr, hosts):
    if len(sys.argv) > 1:
        try:
            i = int(sys.argv[1])
            if i in range(len(hosts)):
                return hosts[i]
        except ValueError:
            if sys.argv[1] in hosts:
                return sys.argv[1]
    scr.clear()
    lines = ['CONFIGURED HOSTS', '']
    lines.extend(to_block('\n'.join(hosts)))
    base_y = (curses.LINES - len(lines) - 1) // 2 + 2
    base_x = (curses.COLS - len(lines[-1]) - 1) // 2
    tip = '<Tab> Move cursor, <Space>/<Enter> Connect, <Esc> Cancel'
    scr.addstr(curses.LINES - 1, (curses.COLS - len(tip) - 1) // 2, tip)
    scr.move(0, 0)
    scr.refresh()
    selected = 0
    while True:
        print_center_core(scr, lines)
        scr.addstr(base_y + selected, base_x, lines[selected + 2], curses.A_REVERSE)
        scr.move(0, 0)
        scr.refresh()
        key = scr.getch()
        if key in (ord('\r'), ord('\n'), ord(' ')):
            return hosts[selected]
        elif key in (ord('\t'), curses.KEY_DOWN, ord('j')):
            selected = (selected + 1) % len(hosts)
        elif key in (curses.KEY_UP, ord('k')):
            selected = (selected - 1) % len(hosts)
        elif key in (ord('\x1b'), ord('q')):
            sys.exit(0)
        elif key in range(ord('1'), ord('9') + 1):
            if key - ord('1') in range(len(hosts)):
                selected = key - ord('1')


def connect_host(scr, host):
    scr.clear()
    scr.move(0, 0)
    curses.nocbreak()
    scr.keypad(False)
    curses.echo()
    curses.endwin()
    return os.system('ssh "{}"'.format(host))


def main(scr):
    hosts = get_hosts()
    if not hosts:
        error_no_host(scr)
        sys.exit(1)
    sys.exit(connect_host(scr, choose_host(scr, hosts)))


if __name__ == '__main__':
    curses.wrapper(main)