import socket, and fix error in ip2host

[python-TSD.git] / lib / TSD / common.py

import ldap
import ldap.sasl
import socket

ldapurl    = 'ldap://tsd-dc01.tsd.usit.no'
ldapbase   = 'DC=tsd,DC=usit,DC=no'

# use 2 to debug ldap connections.
trace_level = 0

def ldap_uri():
    return ldapurl

def ldap_base():
    return ldapbase

def ldap_connect():
    lref = ldap.initialize(ldap_uri(), trace_level=trace_level)

    auth_tokens = ldap.sasl.gssapi()
    lref.sasl_interactive_bind_s('', auth_tokens)
    return lref

def ldap_unbind(lref):
    lref.unbind()

def computermatch(enabled = None):
    enablematch = '(userAccountControl:1.2.840.113556.1.4.803:=2)'
    if enabled is None:
        match = '(objectclass=computer)'
    elif enabled:
        match = '(&(objectclass=computer)(!%s))' % enablematch
    else:
        match = '(&(objectclass=computer)%s)' % enablematch
    return match

def role2hosts(lref, rolename, enabled = None):
    """
Look up all hosts with a given role.
    """
    hosts = {}
    policybase = 'OU=hostpolicies,' + ldap_base()
    rolebase   = 'OU=roles,' + policybase
    atombase   = 'OU=atoms,' + policybase

    # First look up host members in role
    roleres = lref.search_s(rolebase, ldap.SCOPE_SUBTREE,
                        '(&(objectclass=group)(name=%s))' % rolename,
                        ['member'])

    # Next, look up FQDN of host members
    match = computermatch(enabled)
    for roledn, roleentry in roleres:
        if 'member' in roleentry:
            for memberdn in roleentry['member']:
                hostres = lref.search_s(memberdn, ldap.SCOPE_BASE, 
                                        match , ['dNSHostName'])
                for hostdn, hostent in hostres:
                    hosts[hostent['dNSHostName'][0]] = memberdn;
    return hosts.keys()

def role4host(lref, hostname):
    """
Look up all roles of a given host
    """
    roles = {}
    policybase = 'OU=hostpolicies,' + ldap_base()
    rolebase   = 'OU=roles,' + policybase

    # First look up hostname
    hostres = lref.search_s(ldap_base(), ldap.SCOPE_SUBTREE, \
                            '(&(objectclass=computer)(dNSHostName=%s))' % hostname.lower(), \
                            ['memberOf'])
    
    # FIXME fail if more than one host is returned
    
    for hostdn, hostentry in hostres:
        if 'memberOf' in hostentry:
            for roledn in hostentry['memberOf']:
                # look up the role name
                entry = lref.search_s(roledn, ldap.SCOPE_BASE, \
                                      '(objectclass=group)', \
                                      ['name'])
                for dn, roleent in entry:
                    roles[roleent['name'][0]] = 1
    return roles.keys()


def atom2hosts(lref, atomname, enabled = None):
    """
Look up all hosts affected by a given cfengine atom (via role)
    """
    hosts = {}
    policybase = 'OU=hostpolicies,' + ldap_base()
    rolebase   = 'OU=roles,' + policybase
    atombase   = 'OU=atoms,' + policybase

    # First look up role members in the atom
    atomres = lref.search_s(atombase, ldap.SCOPE_SUBTREE,
                        '(&(objectclass=group)(name=%s))' % atomname,
                        ['member'])

    # Next, look up FQDN of host members
    match = computermatch(enabled)
    for atomdn, atomentry in atomres:
        if 'member' in atomentry:
            for roledn in atomentry['member']:
                roleres = lref.search_s(roledn, ldap.SCOPE_BASE, 
                                        '(objectclass=group)',
                                        ['member'])
                for roledn, roleentry in roleres:
                    if 'member' in roleentry:
                        for memberdn in roleentry['member']:
                            hostres = lref.search_s(memberdn, ldap.SCOPE_BASE, 
                                                    match, ['dNSHostName'])
                            for hostdn, hostent in hostres:
                                hosts[hostent['dNSHostName'][0]] = memberdn;
    return hosts.keys()

def host2project(lref, fqdn):
    """
Find computer object and its project based on FQDN.
    """
    res = lref.search_s(ldap_base(), ldap.SCOPE_SUBTREE, \
                        '(&(objectclass=computer)(dnshostname=%s))' % fqdn,
                        ['dn'])
    for dn, entry in res:
        if None is not dn:
            projectdn = ",".join((dn.split(","))[2:])
            projectnum = (dn.split(","))[-5].replace("OU=p", "")
            return (projectnum, dn)
    return (None, None)

def ip2host(lref, hostname):
    """
Find network information about a host
    """
    network = {}
    dnsinfo = socket.getaddrinfo(hostname, None)
    for (family, socktype, proto, canonname, sockaddr) in dnsinfo:
        if socket.AF_INET == family:
            network['ipv4addr'] = sockaddr[0]
        elif socket.AF_INET6 == family:
            network['ipv6addr'] = sockaddr[0]
    hostres = lref.search_s(ldap_base(), ldap.SCOPE_SUBTREE, \
                            '(&(objectclass=computer)(dNSHostName=%s))' % hostname.lower(), \
                            ['networkAddress'])
    for hostdn, hostaddr in hostres:
        ipv6 = ""
        ipv4 = ""
        if 'networkAddress' in hostaddr:
            for i in hostaddr['networkAddress']:
                if ":" in i:
                    ipv6 = i
                elif "." in i:
                    ipv4 = i
                else:
                    network['vlan'] = i
            if "/" in ipv4:
                vlan4, network4, prefix4 = ipv4.split("/")
                network['vlan'] = vlan4
                network['ipv4prefix'] = prefix4
                network['ipv4network'] = network4
            if "/" in ipv6:
                vlan6, network6, prefix6 = ipv6.split("/")
                network['ipv6prefix'] = prefix6
                network['ipv6network'] = network6
                network['vlan'] = vlan6
    return network

def testfuncs():
    l = ldap_connect()
    print "All hosts with role autoinstall_virtual_host"
    print role2hosts(l, 'autoinstall_virtual_host');
    print role2hosts(l, 'autoinstall_virtual_host', enabled=True);
    print "all roles of host tsd-prism.tsd.usit.no"
    print role4host(l, 'tsd-prism.tsd.usit.no')
    print "project number for host tsd-prism.tsd.usit.no"
    print host2project(l, 'tsd-prism.tsd.usit.no')
    print atom2hosts(l, 'root_for_dba')
    print atom2hosts(l, 'root_for_dba', enabled=False)
    ldap_unbind(l)

if __name__ == "__main__":
    testfuncs()