jrtechs-bash_manager/src/ssh_manager.py

"""
Jeffery Russell
9-26-17
"""

import subprocess
import collections

from utils import module
import configuration
import mount_ssh_drive
import os

INPUT_FILE = configuration.get_config()["servers"]
FORWARD_FILE = configuration.get_config()["forwards"]
Computer = collections.namedtuple("Computer", ('host', 'menu_id'))

WELCOME_MESSAGE = "**************************************"


def main():
    """
    This function inputs all the available hosts from a text file and 
    prompts the user to connect to them
    :return:
    """
    file = module.input_file(INPUT_FILE)

    cmp = []
    count = 1

    for line in file:
        cmp.append(Computer(line, count))
        count += 1

    menu = []

    for c in cmp:
        menu.append(str(c.menu_id) + ") " + c.host)
    menu.append("A) Exit")
    menu.append("B) Manager tools")
    menu.append("C) Socks Tunnel")
    menu.append("D) SSH Drive Manager")
    menu.append("E) Local Port Forwarding")

    module.print_menu("SSH manager V 1.0", menu)

    i = input("Enter Option:")

    if i == '' or i == 'A' or i == 'a':
        exit_program()
    elif i == 'b' or i == 'B':
        sub_menu()
    elif "c" == str.lower(i):
        socks_ssh_tunnel()
    elif "d" == str.lower(i):
        mount_ssh_drive.main()
    elif "e" == str.lower(i):
        port_forwarding_sub_menu()
    else:
        for c in cmp:
            if int(i) == c.menu_id:
                subprocess.call(["ssh", c.host])
                exit_program()


def socks_ssh_tunnel():
    """
    prints user a menu to select a host to start a socks proxy with
    :return: None
    """
    file = module.input_file(INPUT_FILE)
    cmp = []
    count = 1
    for line in file:
        cmp.append(Computer(line, count))
        count += 1

    menu = []
    for c in cmp:
        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 'a' == str.lower(i):
        exit_program()
    elif 'b' == str.lower(i):
        main()
    else:
        for c in cmp:
            if int(i) == c.menu_id:
                print_red("Starting socks proxy on " + c.host + ":8123")
                subprocess.call(["ssh", "-D", "8123", "-C", "-q", "-N", c.host])
                exit_program()


def print_sub_menu():
    """
    prints out a sub help menu for other options
    :return: None
    """
    module.print_menu("Options", ["1) Add Host",
                                "2) Copy SSH key to server",
                                "3) Remove host name",
                                "4) Return to ssh manager",
                                "5) Manage Configuration and Bash",
                                "6) Exit"])


def print_red(prt): return "\033[91m {}\033[00m" .format(prt)


def sub_menu():
    """
    calls printSubMenu and then gets input from user to 
    make appropriate function calls
    :return: None
    """
    print_sub_menu()
    i = input("Enter selection:")

    if i != '' and int(i) in {1, 2, 3, 4, 5, 6}:
        options = {1: add_host,
                   2: copy_ssh_key,
                   3: remove_host,
                   4: main,
                   5: configuration.main,
                   6: exit_program
                   }
        options[int(i)]()
    else:
        print("Invalid selection!")

    sub_menu()


def exit_program():
    """
    Exits the program and clears the screen
    :return: None
    """
    subprocess.call(["clear"])
    exit()


def add_host():
    """
    appends an inputted host name to servers.txt
    :return: None
    """
    host = input("Enter 'user@host' or -1 to exit:")
    if host != '-1':
        module.append_file(INPUT_FILE, host)


def copy_ssh_key():
    """
    calls systems ssh-copy-id with host name
    :return: None
    """
    file = module.input_file(INPUT_FILE)
    cmp = []
    count = 1
    for line in file:
        cmp.append(Computer(line, count))
        count += 1

    menu = []
    for c in cmp:
        menu.append(str(c.menu_id) + ") " + c.host)
    menu.append("A) Exit")
    module.print_menu("Copy SSH Key", menu)

    host_id = input("Enter number of host to copy ssh key to:")
    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])


def remove_host():
    """
    Removes a host name from servers.txt
    :return: None
    """
    file = module.input_file(INPUT_FILE)
    cmp = []
    count = 1
    print(print_red("*" * len(WELCOME_MESSAGE)))
    for line in file:
        cmp.append(Computer(line, count))
        count += 1
    for c in cmp:
        space = " " * (len(WELCOME_MESSAGE) - 3 - 
                len(str(c.menu_id) + ") " + c.host))
        print(print_red("*") + str(c.menu_id) + ") " + 
            c.host + space + print_red("*"))

    print(print_red("*" * len(WELCOME_MESSAGE)))

    host = input("Enter number of host -1 to exit:")
    if host != '-1':
        for c in cmp:
            if c.menu_id == int(host):
                module.remove_line_from_file(INPUT_FILE, c.host)

def add_forward_config():
    """
    appends new local forward configurations to forwards.txt
    :return: None
    """
    forwardingConfig = input("Enter forwarding config to add or -1 to exit:")
    if forwardingConfig != '-1':
        module.append_file(FORWARD_FILE, forwardingConfig)


def del_forward_config():
    """
    removes a local forward configurations from forwards.txt
    :return: None
    """
    # check if file is empty before trying to open to read file
    if (os.stat(FORWARD_FILE).st_size == 0):
        print(print_red("*" * len(WELCOME_MESSAGE)))
        print("No forward configurations are available for removal.")
        print(print_red("*" * len(WELCOME_MESSAGE)))

    else:
        f = open(FORWARD_FILE, "r")
        i = 1
        savedConfigs = []
        print(print_red("*" * len(WELCOME_MESSAGE)))
        print("Current saved local port forwarding configurations:")
        for line in f:
            print(print_red("*") + "%s) %s" % (i, line) + print_red("*"))
            savedConfigs.append(line.strip())
            i += 1
        print(print_red("*" * len(WELCOME_MESSAGE)))
        selection = input("Select a configuration to delete: ")
        module.remove_line_from_file(FORWARD_FILE, savedConfigs[int(selection)-1])

def select_forward_config():
    """
    Select a local forwarding configuration from forwards.txt
    to execute using SSH
    :return: None
    """
    # check if file is empty before trying to open to read file
    if (os.stat(FORWARD_FILE).st_size == 0):
        print(print_red("*" * len(WELCOME_MESSAGE)))
        print("No forward configurations are available to run.")
        print(print_red("*" * len(WELCOME_MESSAGE)))

    else:
        f = open(FORWARD_FILE, "r")
        i = 1
        savedConfigs = []
        print(print_red("*" * len(WELCOME_MESSAGE)))
        print("Current saved local port forwarding configurations:")
        for line in f:
            print(print_red("*") + "%s) %s" % (i, line) + print_red("*"))
            savedConfigs.append(line.strip())
            i += 1
        print(print_red("*" * len(WELCOME_MESSAGE)))
        selection = input("Select a configuration to execute: ")

        #Split the selected entry based on a space so we can feed each
        #element into the subprocess call
        forwardConfig = savedConfigs[int(selection)-1].split(" ")
        subprocess.call(["ssh", "-tt", "-L", forwardConfig[0], forwardConfig[1]])
        exit_program()

def print_forwarding_sub_menu():
    """
    prints out a sub help menu for other forwarding options  
    :return: None
    """
    module.print_menu("Options", ["1) Add local forward configuration",
                                  "2) Remove local forward configuration",
                                  "3) Select local forward configuration",
                                  "4) Exit"])

def port_forwarding_sub_menu():
    """

    print out a sub menu related to port forwarding
    options
    :return: None
    """
    print_forwarding_sub_menu()

    i = input("Enter selection:")
    if i != "" and (int(i) in {1, 2, 3, 4}):
        options = {1: add_forward_config,
                   2: del_forward_config,
                   3: select_forward_config,
                   4: exit_program
                   }
        options[int(i)]()
    else:
        print("Invalid selection!") 

    port_forwarding_sub_menu()


"""
Makes sure that other programs don't execute the main
"""
if __name__ == '__main__':
    try:
        main()
    except KeyboardInterrupt:
        exit_program()