Merge pull request #525 from bwelling/tsig

Adds support for a TSIG record class.

2 files changed, 50 insertions, 6 deletions

diff --git a/tests/test_renderer.py b/tests/test_renderer.py
index 345ef82..c60ccf9 100644
--- a/tests/test_renderer.py
+++ b/tests/test_renderer.py
@@ -3,9 +3,11 @@
 import unittest
 
 import dns.exception
+import dns.flags
 import dns.message
 import dns.renderer
-import dns.flags
+import dns.tsig
+import dns.tsigkeyring
 
 basic_answer = \
     """flags QR
@@ -35,6 +37,50 @@ class RendererTestCase(unittest.TestCase):
         expected.id = message.id
         self.assertEqual(message, expected)
 
+    def test_tsig(self):
+        r = dns.renderer.Renderer(flags=dns.flags.RD, max_size=512)
+        qname = dns.name.from_text('foo.example')
+        r.add_question(qname, dns.rdatatype.A)
+        keyring = dns.tsigkeyring.from_text({'key' : '12345678'})
+        keyname = next(iter(keyring))
+        r.write_header()
+        r.add_tsig(keyname, keyring[keyname], 300, r.id, 0, b'', b'',
+                   dns.tsig.HMAC_SHA256)
+        wire = r.get_wire()
+        message = dns.message.from_wire(wire, keyring=keyring)
+        expected = dns.message.make_query(qname, dns.rdatatype.A)
+        expected.id = message.id
+        self.assertEqual(message, expected)
+
+    def test_multi_tsig(self):
+        qname = dns.name.from_text('foo.example')
+        keyring = dns.tsigkeyring.from_text({'key' : '12345678'})
+        keyname = next(iter(keyring))
+
+        r = dns.renderer.Renderer(flags=dns.flags.RD, max_size=512)
+        r.add_question(qname, dns.rdatatype.A)
+        r.write_header()
+        ctx = r.add_multi_tsig(None, keyname, keyring[keyname], 300, r.id, 0,
+                               b'', b'', dns.tsig.HMAC_SHA256)
+        wire = r.get_wire()
+        message = dns.message.from_wire(wire, keyring=keyring, multi=True)
+        expected = dns.message.make_query(qname, dns.rdatatype.A)
+        expected.id = message.id
+        self.assertEqual(message, expected)
+
+        r = dns.renderer.Renderer(flags=dns.flags.RD, max_size=512)
+        r.add_question(qname, dns.rdatatype.A)
+        r.write_header()
+        ctx = r.add_multi_tsig(ctx, keyname, keyring[keyname], 300, r.id, 0,
+                               b'', b'', dns.tsig.HMAC_SHA256)
+        wire = r.get_wire()
+        message = dns.message.from_wire(wire, keyring=keyring,
+                                        tsig_ctx=message.tsig_ctx, multi=True)
+        expected = dns.message.make_query(qname, dns.rdatatype.A)
+        expected.id = message.id
+        self.assertEqual(message, expected)
+
+
     def test_going_backwards_fails(self):
         r = dns.renderer.Renderer(flags=dns.flags.QR, max_size=512)
         qname = dns.name.from_text('foo.example')
diff --git a/tests/test_tsig.py b/tests/test_tsig.py
index 037d5aa..2722e15 100644
--- a/tests/test_tsig.py
+++ b/tests/test_tsig.py
@@ -42,12 +42,11 @@ class TSIGTestCase(unittest.TestCase):
         # not raising is passing
         dns.message.from_wire(w, keyring)
 
-    def make_message_pair(self, qname='example', rdtype='A'):
+    def make_message_pair(self, qname='example', rdtype='A', tsig_error=0):
         q = dns.message.make_query(qname, rdtype)
         q.use_tsig(keyring=keyring, keyname=keyname)
-        q.had_tsig = True  # so make_response() does the right thing
         q.to_wire()  # to set q.mac
-        r = dns.message.make_response(q)
+        r = dns.message.make_response(q, tsig_error=tsig_error)
         return(q, r)
 
     def test_peer_errors(self):
@@ -58,8 +57,7 @@ class TSIGTestCase(unittest.TestCase):
                  (99, dns.tsig.PeerError),
                  ]
         for err, ex in items:
-            q, r = self.make_message_pair()
-            r.tsig_error = err
+            q, r = self.make_message_pair(tsig_error=err)
             w = r.to_wire()
             def bad():
                 dns.message.from_wire(w, keyring=keyring, request_mac=q.mac)