#!/usr/bin/python
# -*- coding: utf-8 -*-

"""
ldap_graph_dot.py

(C) by Michael Stroeder <michael@stroeder.com>

This software is distributed under the terms of the
GPL (GNU GENERAL PUBLIC LICENSE) Version 2
(see http://www.gnu.org/copyleft/gpl.html)

$Id: ldap_graph_dot.py,v 1.3 2011/06/24 17:56:37 michael Exp $

This script generates a DOT file with graph data of LDAP entries found.
"""


import sys,pprint,ldap,ldap.schema,pydot

from ldap.ldapobject import LDAPObject
from ldapurl import LDAPUrl

class MyLDAPUrl(LDAPUrl):
  attr2extype = {
    'who':'bindname',
    'cred':'X-BINDPW',
    'trace_level':'trace',
    'max_level':'maxlevel',
  }

def update_tree_graph(g,l,search_base,max_level):
  if max_level:
    try:
      ldap_result = l.search_s(
        search_base,
        ldap.SCOPE_ONELEVEL,
        filterstr=ldap_url.filterstr or '(objectClass=*)',
        attrlist=ldap_url.attrs or ['objectClass']
      )
    except ldap.NO_SUCH_OBJECT:
      pass
    else:
      for dn,entry in ldap_result:
        if dn:
          struct_oc = schema.get_structural_oc(entry['objectClass'])
          struct_oc_obj = schema.sed[ldap.schema.ObjectClass][struct_oc]
          subordinate_node = pydot.Node(dn)
          subordinate_node.set_shape('box')
          subordinate_node.set_style('rounded')
          subordinate_node.set_label(
            '%s\\n(%s)' % (
              ldap.explode_dn(dn)[0],
              struct_oc_obj.names[0]
            )
          )
          g.add_node(subordinate_node)
          g.add_edge(pydot.Edge(search_base,dn))
          update_tree_graph(g,l,dn,max_level-1)
  return # update_tree_graph()


ldap_url = MyLDAPUrl(sys.argv[1])

ldap_trace_level = int(ldap_url.trace_level or '0')
ldap.trace_level = ldap_trace_level

l = LDAPObject(ldap_url.initializeUrl(),trace_level=ldap_trace_level)
l.protocol_version = 3
l.set_option(ldap.OPT_REFERRALS,0)
l.simple_bind_s((ldap_url.who or ''),(ldap_url.cred or ''))

subschema_dn = l.search_subschemasubentry_s(ldap_url.dn)
schema = ldap.schema.subentry.SubSchema(l.read_subschemasubentry_s(subschema_dn))

tree_graph = pydot.Dot(graph_name='LDAP_TREE',type='di')

tree_graph.add_node(pydot.Node(ldap_url.dn))

update_tree_graph(tree_graph,l,ldap_url.dn,int(ldap_url.max_level or '20'))

sys.stdout.write(tree_graph.to_string())