Return code even if bare username

1 files changed, 33 insertions, 8 deletions

diff --git a/src/leap/mx/alias_resolver.py b/src/leap/mx/alias_resolver.py
index 1d478c7..08ebb73 100644
--- a/src/leap/mx/alias_resolver.py
+++ b/src/leap/mx/alias_resolver.py
@@ -19,6 +19,8 @@
 """
 Classes for resolving postfix aliases.
 
+Test this with postmap -v -q "foo" tcp:localhost:4242
+
 TODO:
     o Look into using twisted.protocols.postfix.policies classes for
       controlling concurrent connections and throttling resource consumption.
@@ -35,13 +37,27 @@ except ImportError:
     print "for instructions on getting required dependencies."
 
 
+class LEAPPostFixTCPMapserver(postfix.PostfixTCPMapServer):
+    def _cbGot(self, value):
+        if value is None:
+            self.sendCode(500, postfix.quote("NOT FOUND SRY"))
+        else:
+            self.sendCode(200, postfix.quote(value))
+
+
 class AliasResolverFactory(postfix.PostfixTCPMapDeferringDictServerFactory):
+
+    protocol = LEAPPostFixTCPMapserver
+
     def __init__(self, couchdb, *args, **kwargs):
         postfix.PostfixTCPMapDeferringDictServerFactory.__init__(
             self, *args, **kwargs)
         self._cdb = couchdb
 
     def _to_str(self, result):
+        """
+        Properly encodes the result string if any.
+        """
         if isinstance(result, unicode):
             result = result.encode("utf8")
         if result is None:
@@ -49,24 +65,33 @@ class AliasResolverFactory(postfix.PostfixTCPMapDeferringDictServerFactory):
         return result
 
     def spit_result(self, result):
+        """
+        Formats the return codes in a postfix friendly format.
+        """
         if result is None:
-            return defer.succeed("500 NO RESULT")
+            return None
         else:
-            return defer.succeed("200")
+            return defer.succeed(result)
 
     def get(self, key):
+        """
+        Looks up the passed key, but only up to the username id of the key.
+
+        At some point we will have to consider the domain part too.
+        """
         try:
             log.msg("Processing key: %s" % (key,))
             if key.find("@") == -1:
-                log.msg("Ignoring key since it's not an email address")
-                return None
-            key = key.split("@")[0]
-            key = key.split("+")[0]
+                # No proper email address, but we need to continue processing
+                # the query so postmap is happy.
+                log.msg("Key it's not an email address")
+            else:
+                key = key.split("@")[0]
+                key = key.split("+")[0]
             log.msg("Final key to query: %s" % (key,))
             d = self._cdb.queryByLoginOrAlias(key)
-            d.addCallback(self._to_str)
-            d.addErrback(log.err)
 
+            d.addCallback(self._to_str)
             d.addCallback(self.spit_result)
             d.addErrback(log.err)
             return d