add_user.py

'''Interactive utility to add users
'''
import datetime as dt
import getpass
import secrets
from argparse import ArgumentParser
from pathlib import Path
import yaml
from passlib.hash import sha512_crypt
import hashlib

HOSTNAME = 'fabricant'

ACCEPTED_GROUPS = {
    1: 'docker',
}

def main():
    parser = ArgumentParser()
    parser.add_argument('--user_file',
                        type=Path,
                        default=Path(f'./{HOSTNAME}_users.yaml'))

    args = parser.parse_args()
    first_name = input('First Name: ')
    last_name = input('Last Name: ')
    email_address = input('Email: ')
    class_flag = (input('Class [y/N]: ').lower() == 'y')
    if class_flag:
        while True:
            class_code = input('Class code (i.e. cse237c): ').lower()
            if class_code != '':
                break
        while True:
            term = input('Class term (i.e. FA24): ').lower()
            if term != '':
                break

    email_digits = int(hashlib.sha512(email_address.encode()).hexdigest(), 16) % 999
    if class_flag:
        username = f'{class_code}_{term}_{first_name[0].lower()}_{last_name.split()[-1].lower()}_{email_digits}'
    else:
        username = f'{first_name[0].lower()}.{last_name.split()[-1].lower()}.{email_digits}'

    with open(args.user_file, 'r', encoding='utf-8') as handle:
        document = yaml.safe_load(handle)
    existing_usernames = [user['username'] for user in document['users']]
    if username in existing_usernames:
        print(f'User {username} already exists!')
        return

    expiration = None
    while not expiration:
        try:
            expiration_day = dt.date.fromisoformat(input('Expiration: '))
            expiration = dt.datetime.combine(expiration_day, dt.time(23, 59, 59))
        except Exception:
            continue
    ssh_keys = []
    while True:
        new_key = input('SSH Key (enter to escape): ')
        if new_key == '':
            if len(ssh_keys) == 0:
                print('SSH pubkey is required!')
                continue
            break
        ssh_keys.append(new_key)
    
    groups = set()
    for gid, gname in ACCEPTED_GROUPS.items():
        print(f'{gid}: {gname}')
    while True:
        new_group = input('Groups to add (enter to escape): ')
        if new_group == '':
            break
        groups.add(ACCEPTED_GROUPS[int(new_group)])
    
    while True:
        password_attempt_1 = getpass.getpass('Password (leave blank for autogenerated password): ')
        if password_attempt_1 == '':
            new_password = secrets.token_urlsafe(25)[:32]
            break
        password_attempt_2 = getpass.getpass('Repeat password: ')
        if password_attempt_2 == password_attempt_1:
            new_password = password_attempt_1
            break
        print('Passwords do not match!')

    password_hash = sha512_crypt.hash(new_password)
    document =  f'  - username: {username}\n'
    document += f'    name: {first_name} {last_name}\n'
    document += '    authorized_keys:\n'
    for key in ssh_keys:
        document += f'      - {key}\n'
    document += f'    expires: "{expiration.strftime("%Y-%m-%d %H:%M:%S")}"\n'
    if len(groups) != 0:
        document += '    groups:\n'
        for group in groups:
            document += f'      - {group}\n'
    document += f'    password: {password_hash}\n'
    with open(args.user_file, 'a', encoding='utf-8') as handle:
        handle.write(document)

    message =  f'username: {username}@{HOSTNAME}.ucsd.edu\n'
    message += '\n'
    message += f'password: {new_password}\n'
    message += '\n'
    message += 'Please allow up to 24 hours for this account to propagate'
    print()
    print('Send the below message to the user:')
    print(message)

if __name__ == '__main__':
    main()