Refractored ssh_manager.py to use common module utilities written for mount_ssh_drive.py

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)

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