#!/usr/bin/python3
#
# This script helps the user quickly enable standard linux tools such as SSH and IPTables
# It also provides a way for the user to add a custom startup script
#
# Author: Dan Bracey -- 6/10/08
# Revision History:
# 6/17/08 -- 1.000 -- debracey -- Initial Creation

import os
import sys

bindir = os.path.abspath(os.path.dirname(os.path.realpath(sys.argv[0])))
libdir = os.path.join(bindir, "..", "python_lib")

sys.path.append(libdir)

import Internet2Lib       # Library functions
import Internet2Consts    # For constants
import sys                # For stdin
import os                 # For standard OS commands
import readline           # Improve the usability of input. Let you use those fancy "arrow" keys
import pwd                # For checking the existence of a user before setting pw. Could cause nasty infinite loop.

htpasswd_file = "/etc/perfsonar/toolkit/psadmin.htpasswd"
htpasswd_exists = os.path.isfile(htpasswd_file)

class menuOption:
    def __init__(self, displayText, executable):
        self.displayText = displayText
        self.executable  = executable
# End menOptions

# Adds a new user to the system
def add_user():
    # We don't want to get in the way here, just let the shell handle this
    done = False

    ans = input("Enter the user whose account you'd like to add. Just hit enter to exit: ").strip()
    if not ans:
        return

    user = ans

    if htpasswd_exists:
        cmd = "htpasswd " + htpasswd_file + " " + user
    else:
        ans = input("Should this user be able to login via SSH? [yes] ").strip().upper()
        if not ans:
            ans = "YES"

        can_login = False
        if ans[0] == "Y":
            can_login = True

        ans = input("Should this user be able to login to the web interface? [yes] ").strip().upper()
        if not ans:
            ans = "YES"

        is_admin = False
        if ans[0] == "Y":
            is_admin = True

        if (can_login):
            cmd = "/usr/sbin/adduser -m"
        else:
            cmd = "/usr/sbin/adduser -M --shell /bin/false"

        if is_admin:
            cmd += " -G psadmin"

        cmd += " " + user

    retVal = os.system(cmd)
    if (retVal != 0):
        print(Internet2Consts.YELLOW + "Error: couldn't add user " + user + Internet2Consts.NORMAL)
        return

    if not htpasswd_exists:
        password(user)

    return
# End add_user

# Deletes a user from the system
def del_user():
    # We don't want to get in the way here, just let the shell handle this
    done = False

    ans = input("Enter the user whose account you'd like to remove. Just hit enter to exit: ").strip()
    if not ans:
        return

    user = ans

    if htpasswd_exists:
        cmd = "htpasswd -D " + htpasswd_file + " " + user
    else:
        delete_home = False

        if os.path.exists(os.path.join("/", "home", user)):
            ans = input("Would you like to delete "+user+"'s home directory? [yes] ").strip().upper()
            if ans and ans[0] == "Y":
                delete_home = True

        cmd = "/usr/sbin/userdel"
        if delete_home:
            cmd += " -r"
        cmd += " "+user

    retVal = os.system(cmd)

    return
# End del_user

# Sets the password for the account
def password(user=None):
    # We don't want to get in the way here, just let the shell handle this
    done = False

    if (user == None):
        ans = input("Enter the user whose password you'd like to change. Just hit enter to exit: ").strip()
        if not ans:
            return

        user = ans
        user_exists = True

        if htpasswd_exists:
            retVal = os.system("grep -q '^" + user + ":' " + htpasswd_file)
            if int(retVal) != 0:
                user_exists = False
        else:
            try:
                user_info = pwd.getpwnam(user)
            except KeyError:
                user_exists = False

        if not user_exists:
            print("User " + Internet2Consts.MAGENTA + user + Internet2Consts.NORMAL + " does not exist.")
            return

    while not done:
        print("Please specify a password for the " + Internet2Consts.MAGENTA + user + Internet2Consts.NORMAL + " account.")

        if htpasswd_exists:
            cmd = "htpasswd " + htpasswd_file + " " + user
        else:
            cmd = "/usr/bin/passwd " + user

        retVal = os.system(cmd)

        if int(retVal) != 0:
            print(Internet2Consts.YELLOW + "Couldn't set the password to what you specified. Please try again." + Internet2Consts.NORMAL)
        else:
            done = True
    # Done setting password
    print("")
    done = False

    return
# End password

# Initializes the menu options
# @return an initialized list of menu options
def initialize():
    optionList = []

    # Add a user
    item = menuOption("Add a new user", add_user)
    optionList.append(item)

    # Delete the user
    item = menuOption("Delete a user", del_user)
    optionList.append(item)

    # Set passord
    item = menuOption("Change a user's password", password)
    optionList.append(item)

    return optionList
# End initialize

# Displays the menu of things to setup
def displayMenu():
    print("You may configure any of the options below with this program: ")

    for i, option in enumerate(menuOptions):
        print("\t" + str(i+1) + ". " + option.displayText)

    print("\t" + "0. exit")

    return
# End displayMenu

# Processes the menu choice (does a do-while until the choice is 0)
# @param fileHandle the fileHandle we are working with (for the startup file)
def mainLoop():
    done = False

    while not done:
        displayMenu()

        # Process the user's input
        try:
            choice = int(input("Make a selection: ").strip())
        except ValueError:
            choice = -1 # Not valid
        print(" ")

        # Negative numbers are invalid -- as are choices > the length of the choice list (Exit is not in the thing)
        if choice < 0 or choice > len(menuOptions):
            print(Internet2Consts.YELLOW + "Invalid selection" + Internet2Consts.NORMAL)
            continue

        # Exit if choice is 0
        elif choice == 0:
            done = True
        else:
            menuOptions[choice-1].executable() # 0 is exit (so needs -1)
            print("")

    return
# End mainLoop

### MAIN ###
print("Welcome to the perfSONAR Toolkit user administration program.")
print("This program will help you administer users.")

# Must be root
if not Internet2Lib.isRoot():
    print(Internet2Consts.YELLOW + "You must run this script as root. Please rerun this script with root permissions." + Internet2Consts.NORMAL)
    sys.exit(1)

menuOptions = initialize()
mainLoop()

sys.exit(0)
### END MAIN
