- IPv6 support added

1 files changed, 41 insertions, 20 deletions

diff --git a/pyxmpp/resolver.py b/pyxmpp/resolver.py
index ddb83ec..986f96d 100644
--- a/pyxmpp/resolver.py
+++ b/pyxmpp/resolver.py
@@ -27,6 +27,7 @@ __docformat__="restructuredtext en"
 
 import re
 import socket
+from socket import AF_UNSPEC, AF_INET, AF_INET6
 import dns.resolver
 import dns.name
 import dns.exception
@@ -34,7 +35,11 @@ import random
 from encodings import idna
 
 service_aliases={"xmpp-server": ("jabber-server","jabber")}
-ip_re=re.compile(r"\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}")
+
+# should match all valid IP addresses, but can pass some false-positives,
+# which are not valid domain names
+ipv4_re=re.compile(r"^\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}$")
+ipv6_re=re.compile(r"^[0-9a-f]{0,4}:[0-9a-f:]{0,29}:([0-9a-f]{0,4}|\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3})$")
 
 def shuffle_srv(records):
     """Randomly reorder SRV records using their weights.
@@ -87,7 +92,7 @@ def reorder_srv(records):
         ret+=shuffle_srv(tmp)
     return ret
 
-def resolve_srv(domain,service,proto="tcp"):
+def resolve_srv(domain,service, proto="tcp"):
     """Resolve service domain to server name and port number using SRV records.
 
     A built-in service alias table will be used to lookup also some obsolete
@@ -119,17 +124,17 @@ def resolve_srv(domain,service,proto="tcp"):
         return [(rr.target.to_text(),rr.port) for rr in reorder_srv(r)]
     return None
 
-def getaddrinfo(host,port,family=0,socktype=socket.SOCK_STREAM,proto=0,allow_cname=True):
+def getaddrinfo(host, port, family= AF_UNSPEC,
+                socktype = socket.SOCK_STREAM, proto = 0, allow_cname = True):
     """Resolve host and port into addrinfo struct.
 
-    Does the same thing as socket.getaddrinfo, but using `pyxmpp.resolver`. This
-    makes it possible to reuse data (A records from the additional section of
-    DNS reply) returned with SRV records lookup done using this module.
+    Does the same thing as socket.getaddrinfo, but using `dns.resolver`.
 
     :Parameters:
         - `host`: service domain name.
         - `port`: service port number or name.
-        - `family`: address family.
+        - `family`: address family (`AF_INET` for IPv4, 
+           `AF_INET6` for IPv6 or `AF_UNSPEC` for either).
         - `socktype`: socket type.
         - `proto`: protocol number or name.
         - `allow_cname`: when False CNAME responses are not allowed.
@@ -150,20 +155,36 @@ def getaddrinfo(host,port,family=0,socktype=socket.SOCK_STREAM,proto=0,allow_cna
         proto=socket.getprotobyname(proto)
     if type(port)!=int:
         port=socket.getservbyname(port,proto)
-    if family not in (0,socket.AF_INET):
-        raise NotImplementedError,"Protocol family other than AF_INET not supported, yet"
-    if ip_re.match(host):
-        return [(socket.AF_INET,socktype,proto,host,(host,port))]
+    if family not in (AF_UNSPEC, AF_INET, AF_INET6):
+        raise NotImplementedError, "Unsupported protocol family."
+    if ipv4_re.match(host) and family in (AF_UNSPEC, AF_INET):
+        return [(AF_INET, socktype, proto, host, (host, port))]
+    if ipv6_re.match(host) and family in (AF_UNSPEC, AF_INET6):
+        return [(AF_INET6, socktype, proto, host, (host, port))]
     host=idna.ToASCII(host)
-    try:
-        r=dns.resolver.query(host, 'A')
-    except dns.exception.DNSException:
-        r=dns.resolver.query(host+".", 'A')
-    if not allow_cname and r.rrset.name!=dns.name.from_text(host):
-        raise ValueError,"Unexpected CNAME record found for %s" % (host,)
-    if r:
-        for rr in r:
-            ret.append((socket.AF_INET,socktype,proto,r.rrset.name,(rr.to_text(),port)))
+    rtypes = []
+    if family in (AF_UNSPEC, AF_INET6):
+        rtypes.append(("AAAA", AF_INET6))
+    if family in (AF_UNSPEC, AF_INET):
+        rtypes.append(("A", AF_INET))
+    exception = None
+    for rtype, rfamily in rtypes:
+        try:
+            try:
+                r=dns.resolver.query(host, rtype)
+            except dns.exception.DNSException:
+                r=dns.resolver.query(host+".", rtype)
+        except dns.exception.DNSException, err:
+            exception = err
+            continue
+        if not allow_cname and r.rrset.name!=dns.name.from_text(host):
+            raise ValueError,"Unexpected CNAME record found for %s" % (host,)
+        if r:
+            for rr in r:
+                ret.append((rfamily, socktype, proto, r.rrset.name,
+                                                        (rr.to_text(),port)))
+    if not ret and exception:
+        raise exception
     return ret
 
 # vi: sts=4 et sw=4