3 files changed, 47 insertions, 30 deletions

diff --git a/src/leap/mx/alias_resolver.py b/src/leap/mx/alias_resolver.py
index 9206ffb..a139dd0 100644
--- a/src/leap/mx/alias_resolver.py
+++ b/src/leap/mx/alias_resolver.py
@@ -30,9 +30,8 @@ TODO:
 
 from twisted.protocols import postfix
 
-from leap.mx.tcp_map import LEAPostfixTCPMapServerFactory
+from leap.mx.tcp_map import LEAPPostfixTCPMapServerFactory
 from leap.mx.tcp_map import TCP_MAP_CODE_SUCCESS
-from leap.mx.tcp_map import TCP_MAP_CODE_TEMPORARY_FAILURE
 from leap.mx.tcp_map import TCP_MAP_CODE_PERMANENT_FAILURE
 
 
@@ -41,30 +40,25 @@ class LEAPPostfixTCPMapAliasServer(postfix.PostfixTCPMapServer):
     A postfix tcp map alias resolver server.
     """
 
-    def _cbGot(self, value):
+    def _cbGot(self, uuid):
         """
         Return a code and message depending on the result of the factory's
         get().
 
-        :param value: The uuid and public key.
+        :param value: The uuid.
         :type value: list
         """
-        uuid, pubkey = value
         if uuid is None:
             self.sendCode(
                 TCP_MAP_CODE_PERMANENT_FAILURE,
                 postfix.quote("NOT FOUND SRY"))
-        elif pubkey is None:
-            self.sendCode(
-                TCP_MAP_CODE_TEMPORARY_FAILURE,
-                postfix.quote("4.7.13 USER ACCOUNT DISABLED"))
         else:
             self.sendCode(
                 TCP_MAP_CODE_SUCCESS,
                 postfix.quote(uuid))
 
 
-class AliasResolverFactory(LEAPostfixTCPMapServerFactory):
+class AliasResolverFactory(LEAPPostfixTCPMapServerFactory):
     """
     A factory for postfix tcp map alias resolver servers.
     """
diff --git a/src/leap/mx/check_recipient_access.py b/src/leap/mx/check_recipient_access.py
index cf172c7..0977564 100644
--- a/src/leap/mx/check_recipient_access.py
+++ b/src/leap/mx/check_recipient_access.py
@@ -23,6 +23,7 @@ Test this with postmap -v -q "foo" tcp:localhost:2244
 """
 
 from twisted.protocols import postfix
+from twisted.internet import defer
 
 from leap.mx.tcp_map import LEAPPostfixTCPMapServerFactory
 from leap.mx.tcp_map import TCP_MAP_CODE_SUCCESS
@@ -33,6 +34,10 @@ from leap.mx.tcp_map import TCP_MAP_CODE_PERMANENT_FAILURE
 class LEAPPostFixTCPMapAccessServer(postfix.PostfixTCPMapServer):
     """
     A postfix tcp map recipient access checker server.
+
+    The server potentially receives the uuid and a PGP key for the user, which
+    are looked up by the factory, and will return a permanent or a temporary
+    failure in case either the user or the key don't exist, respectivelly.
     """
 
     def _cbGot(self, value):
@@ -61,5 +66,43 @@ class LEAPPostFixTCPMapAccessServer(postfix.PostfixTCPMapServer):
 
 
 class CheckRecipientAccessFactory(LEAPPostfixTCPMapServerFactory):
+    """
+    A factory for the recipient access checker.
+
+    When queried, the factory looks up the user's uuid and a PGP key for that
+    user and returns the result to the server's _cbGot() method.
+    """
 
     protocol = LEAPPostFixTCPMapAccessServer
+
+    def _getPubKey(self, uuid):
+        """
+        Look up PGP public key based on user uid.
+
+        :param uuid: The user uid.
+        :type uuid: str
+
+        :return: A deferred that is fired with the uuid and the public key, if
+                 available.
+        :rtype: DeferredList
+        """
+        if uuid is None:
+            return defer.succeed([None, None])
+        # properly encode uuid, otherwise twisted complains when replying
+        if isinstance(uuid, unicode):
+            uuid = uuid.encode("utf8")
+        return defer.gatherResults([
+            defer.succeed(uuid),
+            self._cdb.getPubKey(uuid),
+        ])
+
+    def get(self, key):
+        """
+        Look up uuid and PGP public key based on key.
+
+        :param key: The lookup key.
+        :type key: str
+        """
+        d = LEAPPostfixTCPMapServerFactory.get(self, key)
+        d.addCallback(self._getPubKey)
+        return d
diff --git a/src/leap/mx/tcp_map.py b/src/leap/mx/tcp_map.py
index b7066ff..b62441f 100644
--- a/src/leap/mx/tcp_map.py
+++ b/src/leap/mx/tcp_map.py
@@ -18,7 +18,6 @@
 
 
 from twisted.python import log
-from twisted.internet import defer
 from twisted.internet.protocol import ServerFactory
 
 
@@ -42,24 +41,6 @@ class LEAPPostfixTCPMapServerFactory(ServerFactory):
         """
         self._cdb = couchdb
 
-    def _getPubKey(self, uuid):
-        """
-        Look up PGP public key based on user uid.
-
-        :param uuid: The user uid.
-        :type uuid: str
-
-        :return: A deferred that is fired with the uuid and the public key, if
-                 available.
-        :rtype: DeferredList
-        """
-        if uuid is None:
-            return defer.succeed([None, None])
-        return defer.gatherResults([
-            defer.succeed(uuid),
-            self._cdb.getPubKey(uuid),
-        ])
-
     def get(self, key):
         """
         Look up uuid based on key, only up to the username id of the key.
@@ -71,6 +52,5 @@ class LEAPPostfixTCPMapServerFactory(ServerFactory):
         """
         log.msg("Query key: %s" % (key,))
         d = self._cdb.queryByAddress(key)
-        d.addCallback(self._getPubKey)
         d.addErrback(log.err)
         return d