mirror of https://github.com/azlux/ssh-prompter
260 lines
9.7 KiB
Python
260 lines
9.7 KiB
Python
#!/usr/bin/env python3
|
|
from os import system as system_call, popen
|
|
from os import execlp
|
|
from pathlib import Path
|
|
import signal
|
|
import argparse
|
|
import re
|
|
import curses
|
|
import filecmp
|
|
from typing import Optional
|
|
|
|
search = ""
|
|
args: Optional[argparse.ArgumentParser] = None
|
|
|
|
|
|
def ctrl_caught(signal, frame):
|
|
print("\nSIGINT caught, quitting")
|
|
system_call('if [ -n "$TMUX" ]; then tmux rename-window -t${TMUX_PANE} $(hostname);fi')
|
|
try:
|
|
curses.endwin()
|
|
except curses.error:
|
|
pass
|
|
exit(1)
|
|
|
|
|
|
def change_name(name):
|
|
system_call('if [ -n "$TMUX" ]; then tmux rename-window -t${TMUX_PANE} "' + name + '";fi; if [ -n "$SECURECRT" ]; then echo -n "\033]2;"' + name + '"\007";fi')
|
|
|
|
|
|
def start_ssh(host, ip="", port=22):
|
|
change_name(host)
|
|
if not args.color:
|
|
execlp('/usr/bin/env', '/usr/bin/env', 'bash', '-c',
|
|
'ssh -oStrictHostKeyChecking=no ' + host + ';if [ -n "$TMUX" ]; then tmux rename-window -t${TMUX_PANE} $(hostname);fi; if [ -n "$SECURECRT" ]; then echo -n "\033]2;$(hostname)\007";fi')
|
|
else:
|
|
execlp('/usr/bin/env', '/usr/bin/env', 'bash', '-c',
|
|
'ssh -oStrictHostKeyChecking=no ' + host + ' | ct ;if [ -n "$TMUX" ]; then tmux rename-window -t${TMUX_PANE} $(hostname);fi; if [ -n "$SECURECRT" ]; then echo -n "\033]2;$(hostname)\007";fi')
|
|
|
|
|
|
def create_table_for_prompt(table):
|
|
new_table = []
|
|
match = re.match(r"d/(\S+)/(.*)", search)
|
|
folder = ""
|
|
search_tp = search
|
|
if match:
|
|
folder = match.groups()[0]
|
|
search_tp = match.groups()[1]
|
|
|
|
for i in table:
|
|
if i[2] and not folder and not search_tp:
|
|
if [str('d/' + i[2]) + '/', ''] not in new_table:
|
|
new_table.append([str('d/' + i[2]) + '/', ''])
|
|
|
|
elif (search_tp and search_tp.lower() in i[0].lower()) or (not search_tp and not folder) or (folder and not search_tp and folder == i[2]):
|
|
i2 = i.copy()
|
|
new_table.append(i2)
|
|
return new_table
|
|
|
|
|
|
def get_ssh_server(path=Path(Path.home() / ".ssh/config")):
|
|
all_ssh = []
|
|
if path.is_file():
|
|
with open(path) as f:
|
|
host = ""
|
|
hostname = ""
|
|
folder = ""
|
|
port = 22
|
|
for line in f:
|
|
line = line.strip()
|
|
if len(line) > 0 and line[0] != "#":
|
|
if 'include' in line.lower() and len(line.split(" ")) == 2:
|
|
path_include = line.split(" ")[1]
|
|
if path_include[0] in ['/', '~']:
|
|
pass
|
|
else:
|
|
path_include = "~/.ssh/" + path_include
|
|
new_ssh = get_ssh_server(path=Path(path_include).expanduser())
|
|
all_ssh = all_ssh + new_ssh
|
|
if "host " in line.lower() and "*" not in line.lower():
|
|
if host:
|
|
all_ssh.append([host, hostname, folder, port])
|
|
host = ""
|
|
hostname = ""
|
|
folder = ""
|
|
port = 22
|
|
if "host " in line.lower() and len(line.split(" ")) == 2 and "*" not in line.lower():
|
|
host = line.split(" ")[1]
|
|
elif host:
|
|
if "hostname" in line.lower() and len(line.split(" ")) == 2:
|
|
hostname = line.split(" ")[1]
|
|
elif "folder" in line.lower() and len(line.split(" ")) == 2:
|
|
folder = line.split(" ")[1]
|
|
elif "port" in line.lower() and len(line.split(" ")) == 2:
|
|
port = int(line.split(" ")[1])
|
|
all_ssh.append([host, hostname, folder, port])
|
|
return all_ssh
|
|
else:
|
|
exit(f"No file found at {path}")
|
|
|
|
|
|
def start_prompter(ssh_table, system_argv=None):
|
|
global search
|
|
previous_search = search
|
|
table = create_table_for_prompt(ssh_table)
|
|
if len(table) == 0:
|
|
start_ssh(system_argv[0])
|
|
|
|
screen = curses.initscr()
|
|
curses.noecho()
|
|
curses.cbreak()
|
|
curses.start_color()
|
|
curses.use_default_colors()
|
|
|
|
curses.init_pair(1, curses.COLOR_CYAN, -1)
|
|
|
|
curses.curs_set(0)
|
|
|
|
height, width = screen.getmaxyx()
|
|
title_position_y = 1
|
|
max_box_width = min(width, 80)
|
|
|
|
box_position_y = title_position_y + 3
|
|
box_x_end = width - (width - max_box_width)
|
|
position_x = max(min(width - box_x_end-10, 50), 0)
|
|
|
|
max_lines = max(height - box_position_y - 3, 5)
|
|
|
|
box_search = curses.newwin(3, box_x_end, title_position_y, position_x)
|
|
box_search.box()
|
|
|
|
box = curses.newwin(max_lines + 2, box_x_end, box_position_y, position_x)
|
|
box.keypad(True)
|
|
box.box()
|
|
|
|
cursor_position = 0
|
|
cursor_top = 0
|
|
all_entries = len(table) - 1
|
|
|
|
while True:
|
|
# screen.erase()
|
|
box.erase()
|
|
box.border(0)
|
|
box_search.erase()
|
|
box_search.border(0)
|
|
box_search.addstr(0, 3, "Server List")
|
|
box_search.addstr(1, 5, "Type to search : {}".format(search))
|
|
# box_search.addstr(2, 5, f"debug: cursor_top:{cursor_top} cursor_position:{cursor_position} max:{max_lines} all:{all_entries}")
|
|
# box_search.addstr(2, 5, f"debug: max:{max_lines} height:{height} width:{width}")
|
|
scroll_nb = int((cursor_position * min(max_lines, all_entries)) / all_entries) if all_entries > 0 else 0
|
|
for nb, val in enumerate(table[cursor_top:cursor_top + max_lines]):
|
|
if nb >= max_lines:
|
|
break
|
|
if nb == cursor_position - cursor_top:
|
|
color = curses.A_REVERSE
|
|
else:
|
|
color = curses.A_NORMAL
|
|
|
|
match_folder = re.match(r"d/(\S+)/", val[0])
|
|
if match_folder:
|
|
box.addstr(nb + 1, 2, f"+ {match_folder.groups()[0]}", color)
|
|
else:
|
|
scroll = ' '
|
|
if all_entries > max_lines:
|
|
if (nb - 1 <= scroll_nb <= nb + 1) or (scroll_nb < 2 and nb < 3) or (scroll_nb > max_lines - 2 and nb > max_lines - 2):
|
|
scroll = '#'
|
|
else:
|
|
scroll = '|'
|
|
number_space = min(80, box_x_end) - len(val[0]) - len(val[1]) - 5
|
|
box.addnstr(nb + 1, 2, f"{val[0]}{' ' * number_space}{val[1]}", box_x_end - 4, color)
|
|
box.addnstr(nb + 1, box_x_end - 2, f"{scroll}", box_x_end - 2)
|
|
|
|
# screen.refresh()
|
|
box.refresh()
|
|
box_search.refresh()
|
|
char = box.getch()
|
|
|
|
if char == curses.KEY_UP:
|
|
if len(table) > 0 and cursor_position > 0:
|
|
if cursor_position <= cursor_top + 5 and cursor_top != 0:
|
|
cursor_top -= 1
|
|
cursor_position -= 1
|
|
else:
|
|
cursor_position -= 1
|
|
elif char == curses.KEY_DOWN:
|
|
if len(table) > 0 and cursor_position < all_entries:
|
|
if cursor_position - cursor_top >= max_lines - 5:
|
|
cursor_top += 1
|
|
cursor_position += 1
|
|
else:
|
|
cursor_position += 1
|
|
elif char == curses.KEY_PPAGE:
|
|
if len(table) > 0 and cursor_position > 0:
|
|
cursor_top -= min(max_lines - 1, cursor_top)
|
|
cursor_position -= min(max_lines - 1, cursor_position)
|
|
|
|
elif char == curses.KEY_NPAGE:
|
|
if len(table) > 0 and cursor_position < all_entries:
|
|
cursor_top += min(max_lines - 1, all_entries - cursor_position)
|
|
cursor_position += min(max_lines - 1, all_entries - cursor_position)
|
|
|
|
elif char == curses.KEY_RIGHT or char == curses.KEY_LEFT:
|
|
pass
|
|
elif char == curses.KEY_ENTER or char == 10:
|
|
if len(table) == 0:
|
|
curses.endwin()
|
|
exit()
|
|
match = re.match(r"d/(\S+)/", table[cursor_position % len(table)][0])
|
|
if match:
|
|
search = table[cursor_position % len(table)][0]
|
|
else:
|
|
break
|
|
elif char == curses.KEY_DC or char == 127 or char == curses.KEY_BACKSPACE:
|
|
search = search[:-1]
|
|
elif char == 27: # ESC
|
|
curses.endwin()
|
|
exit()
|
|
else:
|
|
search += chr(char)
|
|
|
|
if previous_search != search:
|
|
table = create_table_for_prompt(ssh_table)
|
|
previous_search = search
|
|
all_entries = len(table) - 1
|
|
cursor_position = 0
|
|
cursor_top = 0
|
|
|
|
curses.endwin()
|
|
|
|
if len(table) == 0:
|
|
exit()
|
|
start_ssh(table[cursor_position][0], table[cursor_position][1], port=int(table[cursor_position][3]))
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(prog="Drop Menu of ~/.ssh/config file", formatter_class=argparse.RawTextHelpFormatter)
|
|
parser.add_argument("--fallback", action="store_true", dest="fallback", required=False, help="USELESS")
|
|
parser.add_argument("--color", action="store_true", dest="color", required=False, help="enable coloring the output")
|
|
global args
|
|
args, system_argv = parser.parse_known_args()
|
|
|
|
ssh_server = get_ssh_server()
|
|
all_host = [i[0].lower() for i in ssh_server]
|
|
if len(system_argv) < 2:
|
|
if len(system_argv) == 1 and system_argv[0].lower() in all_host:
|
|
start_ssh(ssh_server[all_host.index(system_argv[0].lower())][0],
|
|
ssh_server[all_host.index(system_argv[0].lower())][1],
|
|
port=ssh_server[all_host.index(system_argv[0].lower())][3])
|
|
elif len(system_argv) == 1:
|
|
global search
|
|
search = system_argv[0]
|
|
else:
|
|
execlp('ssh', 'ssh', *system_argv)
|
|
|
|
start_prompter(ssh_server, system_argv)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
signal.signal(signal.SIGINT, ctrl_caught)
|
|
main()
|