sql-remctl

#!/usr/bin/env python
"""
sql.mit.edu account management system

This module contains convenience methods for user account
manipulation, including but not limited account creation, and password
generation.
"""

import json
import ldap
import ldap.filter
import os
import random
import string
import sys

from sqlalchemy.orm.exc import MultipleResultsFound, NoResultFound
from sqlalchemy.exc import IntegrityError

import database

class NotifyUserError(Exception):
    """
    All errors of this class will be reported through the API as
    having status = '1', meaning that it was an error condition due to
    to a check in the sql-remctl progrma itself. All other Exceptions
    will be reported as status = '2', meaning they are most likely
    programmer error and need to be reported back to the developers.
    """
    pass

def format_response(inp, status=0):
    """
    Format the dictionary representing the response into a valid,
    machine-readable format. This includes injecting a "status" field.
    Currently, the machine-readable format is json, but this may
    change in the future.
    """
    if not 'status' in inp:
        # Don't inject if the status is already there
        inp['status'] = status
    return json.dumps(inp)

def generate_password(length=10):
    """
    Generate a random password of the specified length (10 by default)
    using the ascii letters, digits, and some special characters.
    """
    available_chars = string.ascii_letters + string.digits + '!@#$%^&*()'
    return ''.join([random.SystemRandom().choice(available_chars) for _ in xrange(length)])

def ensure_authorized(original_function):
    """
    Decorator to ensure that the user is authorized before performing
    actions. Add this to any method that performs a sensitive
    operation and it will only be executed if the user is properly
    authorized.
    """
    def ensure_inner(*args, **kwargs):
        username, target = args[0], args[1]
        if not is_authorized(username, target):
            raise NotifyUserError("User '%s' not authorized for '%s'" %
                                  (username, target))
        return original_function(*args, **kwargs)
    return ensure_inner

def get_user_information(target):
    """
    Synchronously query the ldap.mit.edu LDAP server for the user's
    full name, or return '' if not found.
    """
    mit_ldap = ldap.initialize('ldap://ldap.mit.edu')
    mit_ldap.simple_bind_s()
    default_user = ['junk', {'cn': ['']}]
    result, = mit_ldap.search_s(
        'ou=users,ou=moira,dc=mit,dc=edu',
        ldap.SCOPE_SUBTREE,
        ldap.filter.filter_format(
            '(&(objectClass=posixAccount)' +
            '(uid=%s))',
            [target])) or [default_user]
    fullname = result[1]['cn'][0]
    email = target + '@mit.edu' if fullname != '' else ''
    return (fullname, email)

def get_user(s, target):
    """
    Convenience method to query the database and retrieve the User
    record. This method will throw an exception if the uniqueness
    constraint on the user is violated. This method will also throw a
    notify-user-error if the user account is not found.
    """
    try:
        user = s.query(database.User).filter_by(Username=target).one()
    except MultipleResultsFound as e:
        raise Exception('Fatal error: username uniqueness constraint was violated.')
    except NoResultFound as e:
        raise NotifyUserError("User '%s' is not signed up for a sql account" % (target,))
    return user

def get_database(s, db_name):
    """
    Convenience method to query the database and retrieve the Database
    record. This method has the same semantics as get_user().
    """
    try:
        db = s.query(database.Database).filter_by(Name=db_name).one()
    except MultipleResultsFound as e:
        raise Exception('Fatal error: database uniqueness constraint was violated.')
    except NoResultFound as e:
        raise NotifyUserError("Database '%s' does not exist" % (db_name,))
    return db

@ensure_authorized
def account_create(username, target, args):
    """
    Create the specified target account, if the username (originator)
    is authorized.
    """
    s = database.get_session()
    password = generate_password()
    fullname, email = get_user_information(target)
    user = database.User(target, password, fullname, email)
    s.add(user)
    s.add(database.UserQuota(user))
    s.add(database.UserStat(user))
    try:
        s.commit()
    except IntegrityError as e:
        raise NotifyUserError("User '%s' already has a sql account!" % (target,))

    result = s.execute(database.CreateUser(target, '%', password))

    return {'password': password}

@ensure_authorized
def account_delete(username, target, args):
    """
    Delete the specified target account, if the username (originator)
    is authorized.
    """
    s = database.get_session()
    user = get_user(s, target)
    s.delete(user)
    s.commit()
    
    result = s.execute(database.DropUser(target, '%'))

    return {}

def whoami(*args):
    """
    Interrogate the server about your identity. Conveniently, also
    learn if you currently have a sql account.
    """
    kerberos_name = os.environ['REMOTE_USER']
    username, _ = string.split(kerberos_name, '@', 2)
    s = database.get_session()
    exists = True
    try:
        s.query(database.User).filter_by(Username=username).one()
    except:
        exists = False
    return {'krb5_princ': kerberos_name, 'username': username, 'exists': exists}

def is_auth(username, target, args):
    """
    Determine if the specified username is authorized on the target,
    and report the value. This does not reveal account existence.
    """
    return {'result': is_authorized(username, target)}

@ensure_authorized
def password_set(username, target, args):
    """
    Set the target's password to the specified value, if the
    originator is authorized.
    """
    s = database.get_session()
    user = get_user(s, target)
    if len(args) != 1:
        raise NotifyUserError("Invalid number of arguments specified")
    new_password = args[0]
    user.set_password(new_password)
    s.commit()
    result = s.execute(database.ChangePassword(target, '%', new_password))
    return {}

@ensure_authorized
def database_create(username, target, args):
    """
    Create a new database, enforcing the account+dbname naming scheme,
    database limit, and byte quotas.
    """
    s = database.get_session()
    user = get_user(s, target)
    if len(args) != 1:
        raise NotifyUserError("Invalid number of arguments specified")

    db_name = args[0]

    # We can only create databases if we're under the limit
    have_dbs = sum([x.database.bEnabled for x in user.databases])
    if have_dbs >= user.quota.nDatabasesHard:
        raise NotifyUserError("Cannot create database '%s'; '%s' has %d but the limit is %d" \
                            % (db_name, target, have_dbs, user.quota.nDatabasesHard))
    # Similarly, if the databases are too large, fail out
    if user.stat.nBytes > user.quota.nBytesHard:
        raise NotifyUserError("Cannot create database '%s'; '%s' has used %d bytes but the limit is %d" \
                            % (db_name, target, user.stat.nBytes, user.quota.nBytesHard))

    full_db_name = '%s+%s' % (target, db_name)
    db = database.Database(full_db_name)
    s.add(db)
    s.add(database.DBOwner(user, db))
    s.add(database.DBQuota(db))
    try:
        s.commit()
    except IntegrityError, e:
        return {'error': "Database '%s' already exists!" % (full_db_name,), 'where': 'metadata', 'status': 1}

    # Create the actual database
    try:
        result = s.execute(database.CreateDatabase(full_db_name))
    except:
        s.delete(db)
        s.commit()
        return {'error': "Database '%s' already exists!" % (full_db_name,), 'where': 'sql', 'status': 1}
    # And grant the user privileges on it
    result = s.execute(database.Grant(full_db_name, target, '%'))
    return {'db_name': full_db_name}

@ensure_authorized
def database_drop(username, target, args):
    """
    Drop the specified database.
    """
    if len(args) != 1:
        raise NotifyUserError("Invalid number of arguments specified")
    db_name = args[0]
    s = database.get_session()
    db = get_database(s, '%s+%s' % (target, db_name))
    s.delete(db)
    s.commit()
    # Delete the actual database
    result = s.execute(database.DropDatabase(db.Name, ignore=True))
    # And revoke the user privileges on it
    result = s.execute(database.Revoke(db.Name, target, '%'))

    return {}

@ensure_authorized
def database_list(username, target, args):
    """
    List the specific database
    """
    s = database.get_session()
    user = get_user(s, target)
    if len(args) != 0:
        raise NotifyUserError("Invalid number of arguments specified")

    db_info = []
    for db in user.databases:
        if db.database.bEnabled != 1:
            # Database is "disabled", which actually means it's
            # deleted, so don't print it
            continue
        db_info.append({'name': db.database.Name, 'size': db.database.nBytes})

    return {'quota': user.quota.nBytesHard, 'databases': db_info}

def password_set_random(username, target, args):
    """
    Generate a new random password for the specified account. This
    method does not require authorization, because it uses
    password_set, which does.
    """
    new_password = generate_password()
    password_set(username, target, [new_password])
    return {'password': new_password}

@ensure_authorized
def profile_get(username, target, args):
    """
    Get the target's profile, if the originator is authorized.
    """
    s = database.get_session()
    user = get_user(s, target)
    return {'fullname': user.Name, 'email': user.Email}

@ensure_authorized
def profile_set(username, target, args):
    """
    Set the target's profile, if the originator is authorized.
    """
    s = database.get_session()
    user = get_user(s, target)
    if len(args) != 1:
        raise NotifyUserError("Invalid number of arguments specified")
    try:
        argsd = json.loads(args[0])
    except:
        raise NotifyUserError("Unable to parse specified profile")
    if type(argsd) is not dict:
        raise NotifyUserError("Profile must be a dictionary")
    if 'fullname' in argsd:
        user.Name = argsd['fullname']
    if 'email' in argsd:
        user.Email = argsd['email']
    s.commit()

    return {}

def is_authorized(username, target):
    # THE RULES:
    # -- a user is authorized for itself
    # -- a user is authorized on lockers they have an 'a' bit on
    # -- the sql maintainer team is authorized on all queries
    # -- all else is unauthorized
    if username == target:
        return True
    if target == 'tester-achernya':
        return True
    return False

def main():
    # Figure out which function we are supposed to run
    argv = sys.argv
    argc = len(argv)
    if argc == 1:
        raise NotifyUserError('No operation specified. Try `remctl sql help`.')
    base = os.path.basename(argv[0])
    mode = argv[1]
    account = {'create': account_create,
               'delete': account_delete,
               'whoami': whoami,
               'is-auth': is_auth,
               }
    password = {'set': password_set,
                'generate': password_set_random,
                }
    database = {'create': database_create,
                'drop': database_drop,
                'list': database_list,
                }
    profile = {'get': profile_get,
               'set': profile_set
               }
    ops = {'account': account, 'password': password, 'database': database, 'profile': profile}
    op = ops.get(base, {}).get(mode, None)
    if op == None:
        raise NotifyUserError("Operation '%s %s' not known. Try `remctl sql help`." % (base, mode,))

    # Now, figure out what the target locker is. It's possible there
    # isn't one, in which case we use the username as the sole argument
    username = whoami()['username']
    target = None
    args = None
    # Horrible special case: we don't actually want whoami to take a target, so append ''
    if base == 'account' and mode == 'whoami':
        argv += ['']
    try:
        target, args = argv[2], argv[3:]
    except:
        raise NotifyUserError('Insufficient arguments specified')

    print format_response(op(username, target, args))

if __name__ == '__main__':
    try:
        main()
        sys.exit(0)
    except NotifyUserError as nue:
        print format_response({'error': str(nue)}, status=1)
    except Exception as e:
        print format_response({'error': str(e)}, status=2)

    sys.exit(1)