ldapscan/lib/LdapTree.py

from os.path import dirname, realpath

import ldap
import pygraphviz as pgv

from jinja2 import Environment, FileSystemLoader

class LdapTree(object):
    def __init__(self, hosturi, binddn, basedn, password, use_gssapi):
        #ldap.set_option(ldap.OPT_DEBUG_LEVEL, 1)
        self._ldap = ldap.initialize(hosturi)
        """
        Setting ldap.OPT_REFERRALS to 0 was neccessary to query a samba4
        active directory... Currently I don't know if it is a good idea
        to keep it generally here.
        """
        self._ldap.set_option(ldap.OPT_REFERRALS, 0)
        if use_gssapi:
            sasl_auth = ldap.sasl.sasl({},'GSSAPI')
            self._ldap.sasl_interactive_bind_s("", sasl_auth)
        else:
            self._ldap.bind(binddn, password, ldap.AUTH_SIMPLE)
        self._basedn = basedn
        self._ldap_result = []

        self._data = None

    def text(self, filename = None):
        """
        Returns a text representing the directory.
        If filename is given it will be written in that file.
        """
	env = Environment(loader=FileSystemLoader(
	    dirname(dirname(realpath(__file__))) + '/templates'))
	template = env.get_template('simple.txt.j2')
        text = template.render(ldaptree=self).encode('utf8')

        if filename:
            with open(filename, "w") as text_file:
                text_file.write(text)
        else:
            return text

    def graph(self, filename = None):
        """
        Returns an svg representing the directory.
        If filename is given it will be written in that file.
        """
        graph = pgv.AGraph(
                directed=True, charset='utf-8', fixedsize='true', ranksep=0.1)

        graph.node_attr.update(
                style='rounded,filled', width='0', height='0', shape='box',
                fillcolor='#E5E5E5', concentrate='true', fontsize='8.0',
                fontname='Arial', margin='0.03')

        graph.edge_attr.update(arrowsize='0.55')

        self._graph(graph, self._basedn)

        graph.layout(prog='dot')
        if filename:
            graph.draw(path=filename, format='svg')
            return None
        else:
            return graph.draw(format='svg')

    def _graph(self, graph, dn):
        """
        Recursive function creating a graphviz graph from the directory.
        """
        result = self.node(dn)
        minlen = thislen = 1
        edge_start = dn

        for entry in (entry[0] for entry in result):
            if entry:
                point = entry + '_p'
                sub = graph.add_subgraph()
                sub.graph_attr['rank'] = 'same'
                sub.add_node(
                        point, shape='circle', fixedsize='true', width='0.04',
                        label='', fillcolor='transparent')
                #sub.add_node(entry, URL='https://www.google.de/')
                sub.add_node(entry)
                graph.add_edge(edge_start, point, arrowhead='none',
                        minlen=str(minlen))
                graph.add_edge(point, entry)
                edge_start = point
                minlen = self._graph(graph, entry)
                thislen += minlen

        return thislen

    @property
    def all(self):
        if self._data == None:
            self._data = {}
            result = self._ldap.search_s(self._basedn, ldap.SCOPE_SUBTREE)
            for entry in result:
                self._data[entry[0]] = entry[1:][0]

        return self._data

    @property
    def dn_tree(self):
        retval = {}
        for d in self.all.keys():
            current = retval
            for k in reversed(d.split(',')):
                try:
                    current = current[k]
                except:
                    current[k] = {}
                    current = current[k]
        return retval

    @property
    def hirarchy(self):
        return self._hirarchy(self.dn_tree)

    def _hirarchy(self, dn, base=[], depth=0):
        """
        Hirarchie generates a flat list where each parent is
        followed by all its childs, so that in a template one
        can simple iterate over it to display the complete tree.
        Recently I learned that "recursive loops" are possible
        within jinja2. So we can alter the template that it
        ensures child displays correctly
        """
	retval = []
	for d in dn.keys():
            base_name = ','.join(reversed(base))
            name = ','.join(reversed(base + [d]))
	    retval.append((depth, base_name, name))
	    retval += self._hirarchy(dn[d], base + [d], depth+1)
	return retval

    def childs(self, dn):
        """
        Recently I learned that "recursive loops" are possible
        within jinja2. So we can alter the template that it
	ensures child displays correctly. So this function should
        return all child dn's to a given dn.
        """
        return [d for d in self.hirarchy if d[1] == dn]

    def node(self, dn):
        if dn in self.all:
            return self.all[dn]
        else:
            return {}