#!/usr/bin/env python
#
import sys
from optparse import OptionParser, OptionGroup
from getpass import getpass
import ldap
import ldap.sasl

# User flags for "userAccountControl"
#taken from metzes myldap-pub.py script
UF_SCRIPT                               = 0x00000001
UF_ACCOUNTDISABLE                       = 0x00000002
UF_00000004                             = 0x00000004
UF_HOMEDIR_REQUIRED                     = 0x00000008

UF_LOCKOUT                              = 0x00000010
UF_PASSWD_NOTREQD                       = 0x00000020
UF_PASSWD_CANT_CHANGE                   = 0x00000040
UF_ENCRYPTED_TEXT_PASSWORD_ALLOWED      = 0x00000080

UF_TEMP_DUPLICATE_ACCOUNT               = 0x00000100
UF_NORMAL_ACCOUNT                       = 0x00000200
UF_00000400                             = 0x00000400
UF_INTERDOMAIN_TRUST_ACCOUNT            = 0x00000800

UF_WORKSTATION_TRUST_ACCOUNT            = 0x00001000
UF_SERVER_TRUST_ACCOUNT                 = 0x00002000
UF_00004000                             = 0x00004000
UF_00008000                             = 0x00008000

UF_DONT_EXPIRE_PASSWD                   = 0x00010000
UF_MNS_LOGON_ACCOUNT                    = 0x00020000
UF_SMARTCARD_REQUIRED                   = 0x00040000
UF_TRUSTED_FOR_DELEGATION               = 0x00080000

UF_NOT_DELEGATED                        = 0x00100000
UF_USE_DES_KEY_ONLY                     = 0x00200000
UF_DONT_REQUIRE_PREAUTH                 = 0x00400000
UF_PASSWORD_EXPIRED                     = 0x00800000

UF_TRUSTED_TO_AUTHENTICATE_FOR_DELEGATION= 0x01000000
UF_NO_AUTH_DATA_REQUIRED                = 0x02000000

class AccountObject:
    """Base class for User, Group and Computer objects"""
    
    attributes = {}

    def printAttributes(self):        
        str = ""
        for a, v_list in self.attributes.iteritems():
            for v in v_list:
                print a + ": " + v

class UserAccount(AccountObject):
    """Class for representing samba3 user account"""

    def __init__ (self, attributes):
        self.attributes = attributes
        
    def __init__ (self, dn, attributes):
        self.attributes = attributes
        self.attributes['dn'] = [dn]

class OpenLdapConnector:
    """Open Ldap main object."""

    ldapConnection = None

    def __init__ (self, uri, starttls=False):
        try:
            self.ldapConnection = ldap.initialize(uri)
            #self.ldapConnection.set_option(ldap.OPT_REFERRALS, 0)
            if starttls == True:
                self.ldapConnection.start_tls_s()
            auth_tokens = ldap.sasl.gssapi()
            self.ldapConnection.sasl_interactive_bind_s("", auth_tokens)
                        
        except ldap.LDAPError, e:
            print e
            sys.exit()

    def search (self, basedn, searchfilter='(objectClass=*)', scope=ldap.SCOPE_SUBTREE, attrlist=None):
        """Performs ldap search

        :param searchou:
        """
        
        try:
            results = self.ldapConnection.search_s(basedn, scope, filterstr=searchfilter, attrlist=attrlist)
            assert len(results) == 1, "More than one account found"
            for result in results:
                (dn, attributes) = result
                if dn != None:
                    acct = UserAccount(dn, attributes)

        except ldap.LDAPError, e:
            print e
            sys.exit()
        
        return acct

    def modify (self, dn, modlist=[]):
        try:
            self.ldapConnection.modify_s(dn, modlist)
        except ldap.LDAPError, e:
            print e
            sys.exit()

    def printAttributes (self, attributes):
        str = ""
        for a, v_list in attributes.iteritems():
                    str = str + "Name: " + a + "\n"
                    for v in v_list:
                        str = str + "  Value: " + v + "\n"
        print str
                
            
class modify_user_account:
    """modify user account"""

    parser = OptionParser()

    def __init__(self):
        self.parser.add_option("-v", "--verbose", action="store_true", dest="verbose", default=False, help="provide verbose information")
        self.parser.add_option("", "--set_password_asadmin", action="store_true", dest="set_password_asadmin", default=False, help="Sets user password (Requires administrative priveleges)")
        self.parser.add_option("", "--set_password_asuser", action="store_true", dest="set_password_asuser", default=False, help="Sets user password")
        self.parser.add_option("", "--enable_account", dest="enable_account", type='int', default=None, help="enables (1) or disables (0) the account")
        self.parser.add_option("-U", "--username", dest="username", default=None, help="username of the user to modify")
    
        group_openldap = OptionGroup(self.parser, "OpenLDAP Related Options", "Options used for OpenLDAP backend operations")
        group_openldap.add_option("", "--openldap_starttls", dest="openldap_starttls", action="store_true", default=False, help="Start TLS request [default: %default]")
        group_openldap.add_option("-b", "--openldap_basedn", dest="openldap_basedn", help="base DN")
        group_openldap.add_option("-H", "--openldap_uri", dest="openldap_uri", help="LDAP Uniform Resource Identifier(s)")
        self.parser.add_option_group(group_openldap)
   
    def mod_user_password (self, new_pwd, old_pwd=None):        
        encoded_new_pwd = ("\"%s\"" % new_pwd).encode('utf-16-le')
	if old_pwd is not None:
            encoded_old_pwd = ("\"%s\"" % old_pwd).encode('utf-16-le')
            pwd_replace = [(ldap.MOD_ADD, 'unicodePwd', encoded_new_pwd), (ldap.MOD_DELETE, 'unicodePwd', encoded_old_pwd)]
        else:
            pwd_replace = [(ldap.MOD_REPLACE, 'unicodePwd', encoded_new_pwd)]
        return pwd_replace

    def disable_user_acct (self, userAcct, enable=False):       
        flag = None

        if bool(enable) == False:
            if (int(userAcct.attributes['userAccountControl'][0]) & UF_ACCOUNTDISABLE) != UF_ACCOUNTDISABLE:
                flag = int(userAcct.attributes['userAccountControl'][0]) | UF_ACCOUNTDISABLE
        else:
            if (int(userAcct.attributes['userAccountControl'][0]) & UF_ACCOUNTDISABLE) == UF_ACCOUNTDISABLE:
                flag = int(userAcct.attributes['userAccountControl'][0]) & ~UF_ACCOUNTDISABLE

        if flag is not None:
            return [(ldap.MOD_REPLACE, 'userAccountControl', "%d" % flag)]

        return None

    def run(self):
        (options, args) = self.parser.parse_args()

        searchfilter = '(&(objectClass=user)(sAMAccountName=%s))' % options.username
        attrlist = ['userAccountControl']
        #connection = OpenLdapConnector(uri=options.openldap_uri, binddn=options.openldap_binddn, passwd=bind_pwd, starttls=options.openldap_starttls)
        connection = OpenLdapConnector(uri=options.openldap_uri, starttls=options.openldap_starttls)
        acct = connection.search(basedn=options.openldap_basedn, searchfilter=searchfilter, attrlist=attrlist)

        mods = []
        if options.set_password_asadmin == True or options.set_password_asuser == True:
            pwd1 = getpass("Please type new password: ")
            pwd2 = getpass("Please type new password again: ")
            assert pwd1 == pwd2, "Passwords do not match"
            if options.set_password_asuser == True:
                pwdold = getpass("Please type old password: ")
                mods += self.mod_user_password(pwd1, pwdold)
            else:
                mods += self.mod_user_password(pwd1)
            

        if options.enable_account is not None:
            enable_acct = self.disable_user_acct(acct, enable=options.enable_account)
            if enable_acct is not None:
               mods += enable_acct
        
	if len(mods) > 0:
            connection.modify (acct.attributes['dn'][0], mods)
        
if __name__ == "__main__":
    modifier = modify_user_account()
    modifier.run()
