import ldap
import pygraphviz as pgv

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 = []

    def text(self, filename = None):
        """
        Returns a text representing the directory.
        If filename is given it will be written in that file.
        """
        if filename:
            with open(filename, "w") as text_file:
                text_file.write(self._text(self._basedn, 0))
        else:
            return self._text(self._basedn, 0)

    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 _text(self, dn, level):
        """
        Recursive function that returns a string representation of the
        directory where each depth is indicated by a dash.
        """
        result = self._ldap.search_s(dn, ldap.SCOPE_ONELEVEL)
        indent = '-' * level
        text = indent + dn + "\n"

        for entry in (entry[0] for entry in result):
            if entry:
                text += self._text(entry, level + 1)

        return text

    def _graph(self, graph, dn):
        """
        Recursive function creating a graphviz graph from the directory.
        """
        result = self._ldap.search_s(dn, ldap.SCOPE_ONELEVEL)
        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)
                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