diff --git a/src/mount_ssh_drive.py b/src/mount_ssh_drive.py index 23ee654..205b9ea 100644 --- a/src/mount_ssh_drive.py +++ b/src/mount_ssh_drive.py @@ -69,7 +69,7 @@ def unmount_all_drives(): def remove_drive(): """ - + Prompts the user and removes a drive from MOUNT_FILE """ options = [] file = module.input_file(MOUNT_FILE) diff --git a/src/ssh_manager.py b/src/ssh_manager.py index 6893ced..da4ccc6 100644 --- a/src/ssh_manager.py +++ b/src/ssh_manager.py @@ -16,23 +16,6 @@ Computer = collections.namedtuple("Computer", ('host', 'menu_id')) WELCOME_MESSAGE = "**************************************" -def print_welcome_message(): - """ - Prints defined greeting message to terminal - :return: None - """ - print(print_magenta(WELCOME_MESSAGE)) - - -def print_menu_option(s): - """ - Prints each host option - :param s: - :return: - """ - space = " " * (len(WELCOME_MESSAGE) - 4 - len(s)) - print(print_magenta("* ") + s + space + print_magenta("*")) - def main(): """ @@ -40,7 +23,6 @@ def main(): prompts the user to connect to them :return: """ - print_welcome_message() file = module.input_file(INPUT_FILE) cmp = [] @@ -50,16 +32,16 @@ def main(): cmp.append(Computer(line, count)) count += 1 - print(print_magenta("*") + " " + - print_green("SSH manager V 0.2") + " " + print_magenta("*")) + menu = [] + for c in cmp: - print_menu_option(str(c.menu_id) + ") " + c.host) + menu.append(str(c.menu_id) + ") " + c.host) + menu.append("A) Exit") + menu.append("B) Manager tools") + menu.append("C) Socks Tunnel") - print_menu_option("A) Exit") - print_menu_option("B) Manager tools") - print_menu_option("C) Socks Tunnel") + module.print_menu("SSH manager V 1.0", menu) - print(print_magenta("*" * len(WELCOME_MESSAGE))) i = input("Enter Option:") if i == '' or i == 'A' or i == 'a': @@ -80,7 +62,6 @@ def socks_ssh_tunnel(): prints user a menu to select a host to start a socks proxy with :return: None """ - print_welcome_message() file = module.input_file(INPUT_FILE) cmp = [] count = 1 @@ -88,17 +69,19 @@ def socks_ssh_tunnel(): cmp.append(Computer(line, count)) count += 1 - print(print_magenta("*") + " " + - print_green("Socks Tunnel") + " " + print_magenta("*")) + menu = [] for c in cmp: - print_menu_option(str(c.menu_id) + ") " + c.host) - print_menu_option("A) Exit") - print_menu_option("B) Main") - print(print_magenta("*" * len(WELCOME_MESSAGE))) + menu.append(str(c.menu_id) + ") " + c.host) + + menu.append("A) Exit") + menu.append("B) Main") + module.print_menu("Socks Tunnel", menu) + + i = input("Enter option:") - if i == '' or 'c' == str.lower(i): + if i == '' or 'a' == str.lower(i): exit_program() - elif 'c' == str.lower(i): + elif 'b' == str.lower(i): main() else: for c in cmp: @@ -114,21 +97,11 @@ def print_sub_menu(): prints out a sub help menu for other options :return: None """ - print(print_magenta("**************************************")) - print(print_magenta("*") + print_green("Options") + - " " + print_magenta("*")) - print_menu_option("1) Add Host") - print_menu_option("2) Copy SSH key to server") - print_menu_option("3) Remove host name") - print_menu_option("4) Return to ssh manager") - print_menu_option("5) Exit") - print(print_magenta("*" * len(WELCOME_MESSAGE))) - - -def print_magenta(prt): return"\033[95m {}\033[00m" .format(prt) - - -def print_green(prt): return "\033[92m {}\033[00m" .format(prt) + module.print_menu("Options", ["1) Add Host", + "2) Copy SSH key to server", + "3) Remove host name", + "4) Return to ssh manager", + "5) Exit"]) def print_red(prt): return "\033[91m {}\033[00m" .format(prt) @@ -181,8 +154,6 @@ def copy_ssh_key(): calls systems ssh-copy-id with host name :return: None """ - - print_welcome_message() file = module.input_file(INPUT_FILE) cmp = [] count = 1 @@ -190,18 +161,17 @@ def copy_ssh_key(): cmp.append(Computer(line, count)) count += 1 - print(print_magenta("*") + " " + - print_green("Copy SSH Key") + " " + print_magenta("*")) + menu = [] for c in cmp: - print_menu_option(str(c.menu_id) + ") " + c.host) + menu.append(str(c.menu_id) + ") " + c.host) + menu.append("A) Exit") + module.print_menu("Copy SSH Key", menu) - print_welcome_message() host_id = input("Enter number of host to copy ssh key to:") - if host_id != '-1': + if not (host_id == '-1' or host_id.lower() == 'a'): for c in cmp: if c.menu_id == int(host_id): subprocess.call(["ssh-copy-id", c.host]) - print("Host not found in list!") def remove_host():