Ticket #815: Ldap.py

import imp
import ldap
import Bcfg2.Options
import Bcfg2.Server.Plugin

SCOPE_MAP = {
    "base" : ldap.SCOPE_BASE,
    "one" : ldap.SCOPE_ONELEVEL,
    "sub" : ldap.SCOPE_SUBTREE,
}

LDAP_QUERIES = []

def register_query(query):
    LDAP_QUERIES.append(query)

class ConfigFile(Bcfg2.Server.Plugin.FileBacked):
    """
    Config file for the Ldap plugin

    The config file cannot be 'parsed' in the traditional sense as we would
    need some serious type checking ugliness to just get the LdapQuery
    subclasses. The alternative would be to have the user create a list with
    a predefined name that contains all queries.
    The approach implemented here is having the user call a registering
    decorator that updates a global variable in this module.
    """
    def __init__(self, filename, fam):
        self.filename = filename
        Bcfg2.Server.Plugin.FileBacked.__init__(self, self.filename)
        fam.AddMonitor(self.filename, self)

    def Index(self):
        """
        Reregisters the queries in the config file
        
        The config will take care of actually registering the queries,
        so we just load it once and don't keep it.
        """
        global LDAP_QUERIES
        LDAP_QUERIES = []
        imp.load_source("ldap_cfg", self.filename)

class Ldap(Bcfg2.Server.Plugin.Plugin, Bcfg2.Server.Plugin.Connector):
    """
    The Ldap plugin allows adding data from an LDAP server to your metadata.
    """
    name = "Ldap"
    version = "$Revision: $"
    experimental = True
    debug_flag = False
    
    def __init__(self, core, datastore):
        Bcfg2.Server.Plugin.Plugin.__init__(self, core, datastore)
        Bcfg2.Server.Plugin.Connector.__init__(self)
        self.config = ConfigFile(self.data + "/config.py", core.fam)
    
    def get_additional_data(self, metadata):
        try:
            data = {}
            self.debug_log("LdapPlugin debug: found queries " + 
                                              str(LDAP_QUERIES))
            for QueryClass in LDAP_QUERIES:
                query = QueryClass(metadata)
                if query.is_applicable(metadata):
                    self.debug_log("LdapPlugin debug: processing query '" +
                                                           query.name + "'")
                    data[query.name] = query.get_result(metadata)
                else:
                    self.debug_log("LdapPlugin debug: query '" + query.name +
                        "' not applicable to host '" + metadata.hostname + "'")
            return data
        except Exception, error_msg:
            if self.debug_flag:
                raise
            else:
                Bcfg2.Server.Plugin.logger.error("LdapPlugin error: " + 
                                                         str(error_msg))
                return {}

class LdapConnection(object):
    """
    Connection to an LDAP server.
    """
    def __init__(self, host = "localhost", port = 389, 
                       binddn = None, bindpw = None):
        self.host = host
        self.port = port
        self.binddn = binddn
        self.bindpw = bindpw
        self.conn = None
        
    def __del__(self):
        if self.conn:
            self.conn.unbind()
    
    def init_conn(self):
        self.conn = ldap.initialize(self.url)
        if self.binddn is not None and self.bindpw is not None:
            self.conn.simple_bind_s(self.binddn, self.bindpw)
    
    def run_query(self, query):
        if not self.conn:
            self.init_conn()
        return self.conn.search_s(
            query.base,
            SCOPE_MAP[query.scope],
            query.filter,
            query.attrs,
        )
    
    @property
    def url(self):
        return "ldap://" + self.host + ":" + str(self.port)
    
class LdapQuery(object):
    """
    Query referencing an LdapConnection and providing several
    methods for query manipulation.
    """
    
    name = "unknown"
    base = ""
    scope = "sub"
    filter = "(objectClass=*)"
    attrs = None
    connection = None
    result = None
    
    def __init__(self, metadata):
        """
        Override this to alter the query before execution.
        
        In most cases, you will do something like
        
        self.filter = "(cn=" + metadata.hostname + ")"
        
        here.
        """
        pass
    
    def __unicode__(self):
        return "LdapQuery:" + self.name
    
    def is_applicable(self, metadata):
        """
        Overrideable method to determine if the query is to be executed for
        the given metadata object.
        Defaults to true.
        """
        return True
    
    def prepare_query(self, metadata):
        """
        Overrideable method to alter the query based on metadata.
        Defaults to doing nothing.
        """
        pass
    
    def process_result(self, metadata):
        """
        Overrideable method to post-process the query result.
        Defaults to returning the unaltered result.
        """
        return self.result
    
    def get_result(self, metadata):
        """
        Method to handle preparing, executing and processing the query.
        """
        if isinstance(self.connection, LdapConnection):
            self.prepare_query(metadata)
            self.result = self.connection.run_query(self)
            self.result = self.process_result(metadata)
            return self.result
        else:
            Bcfg2.Server.Plugin.logger.error("LdapPlugin error: " + 
              "No valid connection defined for query " + str(self))
            return None

class LdapSubQuery(LdapQuery):
    """
    SubQueries are meant for internal use only and are not added
    to the metadata object. They are useful for situations where
    you need to run more than one query to obtain some data.
    """
    def prepare_query(self, metadata, subquery_data = None):
        """
        Overrideable method to alter the query based on metadata.
        Defaults to doing nothing.
        """
        pass
    
    def get_result(self, metadata, subquery_data = None):
        """
        Method to handle preparing, executing and processing the query.
        """
        if isinstance(self.connection, LdapConnection):
            self.prepare_query(metadata, subquery_data)
            self.result = self.connection.run_query(self)
            self.process_result(metadata, subquery_data)
            return self.result
        else:
            Bcfg2.Server.Plugin.logger.error("LdapPlugin error: " + 
              "No valid connection defined for query " + str(self))
            return None