#!/usr/bin/env python3

# this script helps manage users on the nodes of a cluster
# it can be invoked by regular users to change their own password on all nodes
# it can also be used by an administrator (a user with sudo authority)
# to add or remove users on all nodes

# Cristian Barbarosie, 2021
# cristian.barbarosie@gmail.com
# https://webpages.ciencias.ulisboa.pt/~cabarbarosie/cluster-ubuntu.html

import sys, string, invoke, fabric, getpass

def usage () :
    print ( 'usage: cluster password' )
    print ( '       cluster add user' )
    print ( '       cluster delete user' )
    sys.exit ( 1 )

if ( len ( sys.argv ) < 2 ) : usage()

remote_nodes = ['beta_1','beta_2','beta_3']  # beta_nodes
all_nodes = ['localhost']                 # alpha_node
for n in remote_nodes : all_nodes.append ( n )

if ( sys.argv[1] == 'password' ) :
    curr_pass = getpass.getpass ( 'Current password : ' )
    new_pass = getpass.getpass ( 'New password : ' )
    verif_pass = getpass.getpass ( 'Please retype new password : ' )
    if new_pass != verif_pass :
        sys.stderr.write ( 'passwords do not match, old password kept\n' )
        sys.exit ( 1 )

    provide_curr_pass = invoke.Responder \
        ( pattern = 'Current password:', response = curr_pass+'\n' )
    provide_new_pass = invoke.Responder \
        ( pattern = 'New password:', response = new_pass+'\n' )
    provide_verif_pass = invoke.Responder \
        ( pattern = 'Retype new password:', response = new_pass+'\n' )

    nb_successful = 0
    for node in all_nodes :
        c = fabric.Connection ( host=node, connect_kwargs = {'password':curr_pass} )
        result = c.run ( 'passwd', asynchronous = True, \
                         watchers = [ provide_curr_pass, provide_new_pass, provide_verif_pass ] )
        result.join()
        res = result.join()
        if ( res.return_code ) :
            sys.stderr.write ( 'error : ' )
            for l in res.tail ( 'stderr' ) : sys.stderr.write ( l )
            sys.exit ( 1 )
        nb_successful += 1

    print ( 'successfully changed password on ', nb_successful, ' nodes' )
    sys.exit ( 0 )

else :
    if ( sys.argv[1] not in ['add','delete' ] ) : usage()
    if ( len ( sys.argv ) == 2 ) : usage()
    if ( sys.argv[2] != 'user' ) : usage()

    forbidden_chars = string.punctuation + string.whitespace
    i = forbidden_chars.find(' ')
    assert ( i >= 0 )
    forbidden_chars_fullname = forbidden_chars[:i] + forbidden_chars[i+1:]

    admin_pass_dict = { 'password': getpass.getpass ( 'your password : ' ) }
    sudo_config = fabric.Config ( overrides = { 'sudo': admin_pass_dict } )

    if ( sys.argv[1] == 'add' ) :

        username = input ( 'username to add : ' )
        for c in username : assert ( c not in forbidden_chars )
        fullname = input ( 'full name : ' )
        for c in username : assert ( c not in forbidden_chars_fullname )
        user_pass = getpass.getpass ( 'initial password for ' + username + ' : ' )

        provide_new_pass = invoke.Responder \
            ( pattern = 'New password:', response = user_pass+'\n' )
        provide_verif_pass = invoke.Responder \
            ( pattern = 'Retype new password:', response = user_pass+'\n' )

        c = fabric.Connection ( host = 'localhost', config = sudo_config, \
                                connect_kwargs = admin_pass_dict          )

        # first, ensure username does not exist
        exists = True
        try : c.run ( 'groups ' + username, hide = 'both' )
        except invoke.exceptions.UnexpectedExit : exists = False
        if ( exists ) :
            print ( 'user', username, 'already exists on local node' )
            sys.exit ( 1 )

        c.sudo ( "useradd --create-home --comment '" + fullname + \
                 "' --shell /usr/bin/bash " + username, hide = 'stderr' )
        c.sudo ( 'mkdir --mode=0750 /sci-data/' + username, hide = 'stderr' )
        c.sudo ( 'chown ' + username + ': /sci-data/' + username, hide = 'stderr' )

        if ( user_pass ) :
            result = c.sudo ( 'passwd ' + username, asynchronous = True, \
                              watchers = [ provide_new_pass, provide_verif_pass ], \
                              hide = 'stderr'                                        )
            result.join()
            res = result.join()
            if ( res.return_code ) :
                sys.stderr.write ( 'error : ' )
                for l in res.tail ( 'stderr' ) : sys.stderr.write ( l )
                sys.exit ( 1 )

        for node in remote_nodes :
            c = fabric.Connection \
                ( host=node, config = sudo_config, connect_kwargs = admin_pass_dict )

            # first, ensure username does not exist
            exists = True
            try : c.run ( 'groups ' + username, hide = 'both' )
            except invoke.exceptions.UnexpectedExit : exists = False
            if ( exists ) :
                print ( 'user', username, 'already exists on', node )
                sys.exit ( 1 )

            c.sudo ( "useradd --no-create-home --comment '" + fullname + \
                     "' --shell /usr/bin/bash " + username, hide = 'stderr' )
            c.sudo ( 'mkdir --mode=0750 /sci-data/' + username, hide = 'stderr' )
            c.sudo ( 'chown ' + username + ': /sci-data/' + username, hide = 'stderr' )

            if ( user_pass ) :
                result = c.sudo ( 'passwd ' + username, asynchronous = True, \
                                  watchers = [ provide_new_pass, provide_verif_pass ], \
                                  hide = 'stderr'                                      )
                result.join()
                res = result.join()
                if ( res.return_code ) :
                    sys.stderr.write ( 'error : ' )
                    for l in res.tail ( 'stderr' ) : sys.stderr.write ( l )
                    sys.exit ( 1 )

        print ( 'successfully added user', username, 'on', len(all_nodes), 'nodes' )

    if ( sys.argv[1] == 'delete' ) :

        username = input ( 'username to delete : ' )
        for c in username : assert ( c not in forbidden_chars )

        for node in remote_nodes :
            c = fabric.Connection ( host = node, config = sudo_config, \
                                    connect_kwargs = admin_pass_dict   )

            # first, ensure username exists
            exists = True
            try : c.run ( 'groups ' + username, hide = 'both' )
            except invoke.exceptions.UnexpectedExit : exists = False
            if ( not exists ) :
                print ( 'no such user', username, 'on', node )
                continue

            c.sudo ( 'userdel  ' + username, hide = 'stderr' )
            c.sudo ( 'rm -rf /sci-data/' + username, hide = 'stderr' )

        c = fabric.Connection ( host = 'localhost', config = sudo_config, \
                                connect_kwargs = admin_pass_dict          )

        # first, ensure username exists
        exists = True
        try : c.run ( 'groups ' + username, hide = 'both' )
        except invoke.exceptions.UnexpectedExit : exists = False
        if ( not exists ) :
            print ( 'no such user', username, 'on local node' )
        else :
            c.sudo ( 'userdel --remove ' + username, hide = 'stderr' )
            c.sudo ( 'rm -rf /sci-data/' + username, hide = 'stderr' )

        print ( 'successfully deleted user', username, 'from all nodes' )