Thursday, May 2, 2013

Java: Active directory paging code


Forgive some of the bad code this was a quick script and a reminder to self on how to do this: (note this code changes all the AD users passwords so be careful)

package common.jndi;

import javax.naming.NamingEnumeration;
import javax.naming.NamingException;

import javax.naming.directory.Attribute;
import javax.naming.directory.Attributes;
import javax.naming.directory.BasicAttribute;
import javax.naming.directory.DirContext;
import javax.naming.directory.InitialDirContext;
import javax.naming.directory.ModificationItem;
import javax.naming.directory.SearchControls;
import javax.naming.directory.SearchResult;

import javax.naming.ldap.Control;
import javax.naming.ldap.InitialLdapContext;
import javax.naming.ldap.PagedResultsControl;
import javax.naming.ldap.PagedResultsResponseControl;

public class ModifyLdapUserScript
    extends LdapBase {
    public static final int PAGE_SIZE = 999;
    private static final String DEFAULT_PASSWORD = "Password1";
    private static final String AD_BASE = "dc=dev,dc=system,dc=uni,dc=edu";

    public ModifyLdapUserScript() {
        super();
    }

    public static void main(String[] args)
        throws Exception {
        String adServer = "adserver";
        String adUser = "cn=Admin,cn=Users," + AD_BASE;
        String adPwd = "secretpassword";

        System.out.println("Starting user modification");
        InitialLdapContext ldapContext = null;
        byte[] cookie = null;
        int runningTotal = 0;
        try {
            long startTime = System.currentTimeMillis();
            ldapContext = connectAD(adServer, adUser, adPwd);

            // Create the search controls
            SearchControls searchCtls = new SearchControls();
            //Specify the attributes to return
            String returnedAtts[] = { "cn" };
            searchCtls.setReturningAttributes(returnedAtts);
            ldapContext.setRequestControls(new Control[] { new PagedResultsControl(PAGE_SIZE, Control.CRITICAL) });
            searchCtls.setSearchScope(SearchControls.SUBTREE_SCOPE);

            //specify the LDAP search filter
            String searchFilter = "(&(cn=*)(objectclass=user))";

            //initialize counter to total the results
            PagedResult totalResults = new PagedResult();

            long innerStartTime = System.currentTimeMillis();
            NamingEnumeration answer = ldapContext.search(AD_BASE, searchFilter, searchCtls);
            totalResults = processPage(ldapContext, answer);
            System.out.println("Processed page 1 [" + totalResults.getResultSize() + "] in time:" +
                               (System.currentTimeMillis() - innerStartTime));
            runningTotal += totalResults.getResultSize();
            int i = 2;
            while (totalResults != null && totalResults.getCookie() != null) {
                innerStartTime = System.currentTimeMillis();
                ldapContext.setRequestControls(new Control[] { new PagedResultsControl(PAGE_SIZE, totalResults.getCookie(), Control.CRITICAL) });
                answer = ldapContext.search(AD_BASE, searchFilter, searchCtls);
                totalResults = processPage(ldapContext, answer);
                System.out.println("Processed page " + i + " [" + totalResults.getResultSize() + "] in time:" +
                                   (System.currentTimeMillis() - innerStartTime));
                runningTotal += totalResults.getResultSize();
                i++;
            }

            System.out.println("Time taken changing users[" + runningTotal + "]:" + (System.currentTimeMillis() - startTime));
        } finally {
            disconnectAD(ldapContext);
        }
        System.out.println("Finished modification");
    }

    private static PagedResult processPage(InitialLdapContext ldapContext, NamingEnumeration answer)
        throws NamingException {
        //Loop through the search results
        String adContent = null;
        PagedResult result = new PagedResult();
        int results = 0;
        while (answer != null && answer.hasMoreElements()) {
            SearchResult sr = answer.next();
            Attributes attrs = sr.getAttributes();
            if (attrs != null) {
                String cn = (String) getAttributeValue(attrs, "cn");
                StringBuffer buffer = new StringBuffer("cn=").append(cn).append(",cn=Users,").append(AD_BASE);
                changePassword(ldapContext, DEFAULT_PASSWORD, buffer.toString());
            }
        }

        Control[] controls = ldapContext.getResponseControls();
        if (controls != null) {
            for (int i = 0; i < controls.length; i++) {
                if (controls[i] instanceof PagedResultsResponseControl) {
                    PagedResultsResponseControl prrc = (PagedResultsResponseControl)controls[i];
                    result.setCookie(prrc.getCookie());
                    result.setResultSize(prrc.getResultSize());
                } else {
                    // Handle other response controls (if any)
                }
            }
        }

        return result;
    }

    public static void changePassword(InitialDirContext context, String password, String username)
        throws NamingException {

        System.out.println("Context Sucessfully Initialized");
        ModificationItem[] mods = new ModificationItem[1];
        Attribute mod0 = new BasicAttribute("userpassword", password);
        mods[0] = new ModificationItem(DirContext.REPLACE_ATTRIBUTE, mod0);
        context.modifyAttributes(username, mods);
    }
}