diff --git a/dkim/__init__.py b/dkim/__init__.py index 518a6c0..427bb48 100644 --- a/dkim/__init__.py +++ b/dkim/__init__.py @@ -356,14 +356,8 @@ def verify(message, logger=None, dnsfunc=get_txt): (base64.b64encode(bodyhash), sig[b'bh'])) return False - # dnstxt wants Unicode - try: - selector = sig[b's'].decode('ascii') - domain = sig[b'd'].decode('ascii') - except UnicodeDecodeError: - return False - name = "%s._domainkey.%s." % (selector, domain) - s = dnsfunc(name).encode('utf-8') + name = sig[b's'] + b"._domainkey." + sig[b'd'] + b"." + s = dnsfunc(name) if not s: return False try: diff --git a/dkim/dns.py b/dkim/dns.py index c158f2b..e14c939 100644 --- a/dkim/dns.py +++ b/dkim/dns.py @@ -49,7 +49,20 @@ def get_txt_pydns(name): # Prefer dnspython if it's there, otherwise use pydns. try: import dns.resolver - get_txt = get_txt_dnspython + _get_txt = get_txt_dnspython except ImportError: import DNS - get_txt = get_txt_pydns + _get_txt = get_txt_pydns + + +def get_txt(name): + """Return a TXT record associated with a DNS name. + + @param name: The bytestring domain name to look up. + """ + # pydns needs Unicode, but DKIM's d= is ASCII (already punycoded). + try: + unicode_name = name.decode('ascii') + except UnicodeDecodeError: + return None + return _get_txt(unicode_name).decode('utf-8') diff --git a/dkim/tests/test_dkim.py b/dkim/tests/test_dkim.py index afe7fbc..ec23e22 100644 --- a/dkim/tests/test_dkim.py +++ b/dkim/tests/test_dkim.py @@ -53,8 +53,12 @@ class TestSignAndVerify(unittest.TestCase): self.key = read_test_data("test.private") def dnsfunc(self, domain): + try: + domain = domain.decode('ascii') + except UnicodeDecodeError: + return None self.assertEqual('test._domainkey.example.com.', domain) - return read_test_data("test.txt").decode('utf-8') + return read_test_data("test.txt") def test_verifies(self): # A message verifies after being signed.