From ec3b370e5135d10e9f6eecc3c548bca834874d81 Mon Sep 17 00:00:00 2001
From: Bob Halley <halley@dnspython.org>
Date: Sun, 3 May 2020 07:23:54 -0700
Subject: [PATCH] Add IDNA codec support to tokenizer and dns.rdata.from_text()

---
 dns/rdata.py        | 11 +++++++++--
 dns/tokenizer.py    | 15 +++++++++++++--
 tests/test_rdata.py | 13 +++++++++++++
 3 files changed, 35 insertions(+), 4 deletions(-)

diff --git a/dns/rdata.py b/dns/rdata.py
index 6396d1bb..dafd4c3b 100644
--- a/dns/rdata.py
+++ b/dns/rdata.py
@@ -392,7 +392,7 @@ def get_rdata_class(rdclass, rdtype):
 
 
 def from_text(rdclass, rdtype, tok, origin=None, relativize=True,
-              relativize_to=None):
+              relativize_to=None, idna_codec=None):
     """Build an rdata object from text format.
 
     This function attempts to dynamically load a class which
@@ -420,11 +420,18 @@ def from_text(rdclass, rdtype, tok, origin=None, relativize=True,
     *relativize_to*, a ``dns.name.Name`` (or ``None``), the origin to use
     when relativizing names.  If not set, the *origin* value will be used.
 
+    *idna_codec*, a ``dns.name.IDNACodec``, specifies the IDNA
+    encoder/decoder to use if a tokenizer needs to be created.  If
+    ``None``, the default IDNA 2003 encoder/decoder is used.  If a
+    tokenizer is not created, then the codec associated with the tokenizer
+    is the one that is used.
+
     Returns an instance of the chosen Rdata subclass.
+
     """
 
     if isinstance(tok, str):
-        tok = dns.tokenizer.Tokenizer(tok)
+        tok = dns.tokenizer.Tokenizer(tok, idna_codec=idna_codec)
     cls = get_rdata_class(rdclass, rdtype)
     if cls != GenericRdata:
         # peek at first token
diff --git a/dns/tokenizer.py b/dns/tokenizer.py
index 8585ccc8..547e030f 100644
--- a/dns/tokenizer.py
+++ b/dns/tokenizer.py
@@ -218,9 +218,13 @@ class Tokenizer(object):
     line_number: The current line number
 
     filename: A filename that will be returned by the where() method.
+
+    idna_codec: A dns.name.IDNACodec, specifies the IDNA
+    encoder/decoder.  If None, the default IDNA 2003
+    encoder/decoder is used.
     """
 
-    def __init__(self, f=sys.stdin, filename=None):
+    def __init__(self, f=sys.stdin, filename=None, idna_codec=None):
         """Initialize a tokenizer instance.
 
         f: The file to tokenize.  The default is sys.stdin.
@@ -229,6 +233,10 @@ class Tokenizer(object):
 
         filename: the name of the filename that the where() method
         will return.
+
+        idna_codec: A dns.name.IDNACodec, specifies the IDNA
+        encoder/decoder.  If None, the default IDNA 2003
+        encoder/decoder is used.
         """
 
         if isinstance(f, str):
@@ -254,6 +262,9 @@ class Tokenizer(object):
         self.delimiters = _DELIMITERS
         self.line_number = 1
         self.filename = filename
+        if idna_codec is None:
+            idna_codec = dns.name.IDNA_2003
+        self.idna_codec = idna_codec
 
     def _get_char(self):
         """Read a character from input.
@@ -565,7 +576,7 @@ class Tokenizer(object):
         """
         if not token.is_identifier():
             raise dns.exception.SyntaxError('expecting an identifier')
-        name = dns.name.from_text(token.value, origin)
+        name = dns.name.from_text(token.value, origin, self.idna_codec)
         return name.choose_relativity(relativize_to or origin, relativize)
 
     def get_name(self, origin=None, relativize=False, relativize_to=None):
diff --git a/tests/test_rdata.py b/tests/test_rdata.py
index c41758c7..55c6c2c0 100644
--- a/tests/test_rdata.py
+++ b/tests/test_rdata.py
@@ -1,3 +1,4 @@
+# -*- coding: utf-8
 # Copyright (C) Dnspython Contributors, see LICENSE for text of ISC license
 
 # Copyright (C) 2006, 2007, 2009-2011 Nominum, Inc.
@@ -17,6 +18,7 @@
 
 import unittest
 
+import dns.name
 import dns.rdata
 import dns.rdataclass
 import dns.rdatatype
@@ -104,5 +106,16 @@ class RdataTestCase(unittest.TestCase):
                                     '"foo\u200b\\123bar"')
         self.assertEqual(str(rdata), '"foo\\226\\128\\139{bar"')
 
+    def test_unicode_idna2003_in_rdata(self):
+        rdata = dns.rdata.from_text(dns.rdataclass.IN, dns.rdatatype.NS,
+                                    "Königsgäßchen")
+        self.assertEqual(str(rdata.target), 'xn--knigsgsschen-lcb0w')
+
+    def test_unicode_idna2008_in_rdata(self):
+        rdata = dns.rdata.from_text(dns.rdataclass.IN, dns.rdatatype.NS,
+                                    "Königsgäßchen",
+                                    idna_codec=dns.name.IDNA_2008)
+        self.assertEqual(str(rdata.target), 'xn--knigsgchen-b4a3dun')
+
 if __name__ == '__main__':
     unittest.main()
-- 
2.47.3