ssh-prompter/ssh.py

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()