LDAP proxy with logs

ldap_proxy.py

#! /usr/bin/env python
# encoding: utf8

import sys
from functools import partial

from ldaptor.protocols import pureldap
from ldaptor.protocols.ldap import ldapserver
from ldaptor.protocols.ldap.ldapclient import LDAPClient, \
    LDAPClientConnectionLostException
from ldaptor.protocols.ldap.ldapconnector import connectToLDAPEndpoint
from ldaptor.protocols.ldap.proxybase import ProxyBase
from twisted.internet import defer, protocol, reactor
from twisted.python import log


def tcp_conn_repr(tcp_transport, reverse=False):
    host = tcp_transport.getHost()
    peer = tcp_transport.getPeer()
    if reverse:
        host, peer = peer, host
    return '{} {}:{} <-> {}:{}'.format(hex(id(tcp_transport)), host.host,
                                       host.port, peer.host, peer.port)


def tcp_lost_conn_repr(tcp_transport):
    return hex(id(tcp_transport)), tcp_transport


class LoggingProxy(ProxyBase):
    """
    A simple example of using `ProxyBase` to log requests and responses.
    """
    def handleProxiedResponse(self, response, request, controls):
        """
        Log the representation of the responses received.
        """
        # TODO: 把日志中的utf8字符串解码以方便查看。
        # log.msg("Request => " + repr(request))
        log.msg("{}: Request => {}".
                format(tcp_conn_repr(self.transport), repr(request)))
        # log.msg("Response => " + repr(response))
        log.msg("{}: Response => {}".
                format(tcp_conn_repr(self.transport), repr(response)))
        return defer.succeed(response)

    def connectionMade(self):
        """
        Establish a connection with an LDAP client.
        """
        try:
            self.transport.setTcpKeepAlive(1)
        except AttributeError:
            pass
        assert self.clientConnector is not None, (
            "You must set the `clientConnector` property on this instance.  "
            "It should be a callable that attempts to connect to a server. "
            "This callable should return a deferred that will fire with a "
            "protocol instance when the connection is complete.")
        d = self.clientConnector()
        d.addCallback(self._connectedToProxiedServer)
        d.addErrback(self._failedToConnectToProxiedServer)
        ldapserver.BaseLDAPServer.connectionMade(self)
        log.msg("C -> P CONN MADE: %s" % tcp_conn_repr(self.transport,
                                                       reverse=True))

    def connectionLost(self, reason):
        if self.client is not None and self.client.connected:
            if not self.unbound:
                self.client.unbind()
                self.unbound = True
            else:
                self.client.transport.loseConnection()
        self.client = None
        ldapserver.BaseLDAPServer.connectionLost(self, reason)
        log.msg("C -> P CONN LOST: {} {}".format(
            *tcp_lost_conn_repr(self.transport)))

# Monkey patch LDAPBindRequest
def ldapBindRequestRepr(self):
    l=[]
    l.append('version={0}'.format(self.version))
    l.append('dn={0}'.format(repr(self.dn)))
    l.append('auth=****')
    if self.tag!=self.__class__.tag:
        l.append('tag={0}'.format(self.tag))
    l.append('sasl={0}'.format(repr(self.sasl)))
    return self.__class__.__name__+'('+', '.join(l)+')'
pureldap.LDAPBindRequest.__repr__ = ldapBindRequestRepr


class MyLDAPClient(LDAPClient):
    """An LDAP client that connect to the proxied server."""
    def connectionMade(self):
        """TCP connection has opened"""
        try:
            self.transport.setTcpKeepAlive(1)
        except AttributeError:
            pass
        finally:
            self.connected = 1
            log.msg('P -> S CONN MADE: %s.' % (tcp_conn_repr(self.transport)))

    def connectionLost(self, reason=protocol.connectionDone):
        """Called when TCP connection has been lost"""
        self.connected = 0
        log.msg("P -> S CONN LOST: {} {}".format(
            *tcp_lost_conn_repr(self.transport)))
        # notify handlers of operations in flight
        while self.onwire:
            k, v = self.onwire.popitem()
            d, _, _, _ = v
            d.errback(reason)

    def _send(self, op):
        if not self.connected:
            log.msg("UNEXPECTED P -> S CONN LOST: {} {}".format(
                hex(id(self.transport)), self.transport))
            raise LDAPClientConnectionLostException()
        msg=pureldap.LDAPMessage(op)
        if self.debug:
            log.msg('P -> S %s' % repr(msg))
        assert not self.onwire.has_key(msg.id)
        return msg

if __name__ == '__main__':
    """
    Demonstration LDAP proxy; listens on localhost:10389 and
    passes all requests to localhost:8081.
    """
    log.startLogging(sys.stderr)
    factory = protocol.ServerFactory()
    proxiedEndpointStr = 'tcp:host=ad.server:port=389'
    use_tls = False
    clientConnector = partial(
        connectToLDAPEndpoint,
        reactor,
        proxiedEndpointStr,
        MyLDAPClient)

    def buildProtocol():
        proto = LoggingProxy()
        proto.clientConnector = clientConnector
        proto.use_tls = use_tls
        # This doubles lines of logs.
        # proto.debug = True
        return proto

    factory.protocol = buildProtocol
    port = 8081 if sys.platform == 'darwin' else 389
    reactor.listenTCP(port, factory)
    reactor.run()