Add support for new TSIG algorithms.

This adds support for the hmac-sha256-128, hmac-sha384-192, and hmac-sha512-256 truncated algorithms.

This also reorders some of the declarations in the TSIG code.

2 files changed, 62 insertions, 26 deletions

diff --git a/dns/tsig.py b/dns/tsig.py
index ab45951..d9d3c24 100644
--- a/dns/tsig.py
+++ b/dns/tsig.py
@@ -72,6 +72,22 @@ class PeerBadTruncation(PeerError):
     """The peer didn't like amount of truncation in the TSIG we sent"""
 
 
+# TSIG Algorithms
+
+HMAC_MD5 = dns.name.from_text("HMAC-MD5.SIG-ALG.REG.INT")
+HMAC_SHA1 = dns.name.from_text("hmac-sha1")
+HMAC_SHA224 = dns.name.from_text("hmac-sha224")
+HMAC_SHA256 = dns.name.from_text("hmac-sha256")
+HMAC_SHA256_128 = dns.name.from_text("hmac-sha256-128")
+HMAC_SHA384 = dns.name.from_text("hmac-sha384")
+HMAC_SHA384_192 = dns.name.from_text("hmac-sha384-192")
+HMAC_SHA512 = dns.name.from_text("hmac-sha512")
+HMAC_SHA512_256 = dns.name.from_text("hmac-sha512-256")
+GSS_TSIG = dns.name.from_text("gss-tsig")
+
+default_algorithm = HMAC_SHA256
+
+
 class GSSTSig:
     """
     GSS-TSIG TSIG implementation.  This uses the GSS-API context established
@@ -139,53 +155,54 @@ class HMACTSig:
     HMAC TSIG implementation.  This uses the HMAC python module to handle the
     sign/verify operations.
     """
+
+    _hashes = {
+        HMAC_SHA1: hashlib.sha1,
+        HMAC_SHA224: hashlib.sha224,
+        HMAC_SHA256: hashlib.sha256,
+        HMAC_SHA256_128: (hashlib.sha256, 128),
+        HMAC_SHA384: hashlib.sha384,
+        HMAC_SHA384_192: (hashlib.sha384, 192),
+        HMAC_SHA512: hashlib.sha512,
+        HMAC_SHA512_256: (hashlib.sha512, 256),
+        HMAC_MD5: hashlib.md5,
+    }
+
     def __init__(self, key, algorithm):
         try:
-            digestmod = _hashes[algorithm]
+            hashinfo = self._hashes[algorithm]
         except KeyError:
             raise NotImplementedError(f"TSIG algorithm {algorithm} " +
                                       "is not supported")
 
         # create the HMAC context
-        self.hmac_context = hmac.new(key, digestmod=digestmod)
+        if isinstance(hashinfo, tuple):
+            self.hmac_context = hmac.new(key, digestmod=hashinfo[0])
+            self.size = hashinfo[1]
+        else:
+            self.hmac_context = hmac.new(key, digestmod=hashinfo)
+            self.size = None
         self.name = self.hmac_context.name
+        if self.size:
+            self.name += f'-{self.size}'
 
     def update(self, data):
         return self.hmac_context.update(data)
 
     def sign(self):
         # defer to the HMAC digest() function for that digestmod
-        return self.hmac_context.digest()
+        digest = self.hmac_context.digest()
+        if self.size:
+            digest = digest[: (self.size // 8)]
+        return digest
 
     def verify(self, expected):
         # re-digest and compare the results
-        mac = self.hmac_context.digest()
+        mac = self.sign()
         if not hmac.compare_digest(mac, expected):
             raise BadSignature
 
 
-# TSIG Algorithms
-
-HMAC_MD5 = dns.name.from_text("HMAC-MD5.SIG-ALG.REG.INT")
-HMAC_SHA1 = dns.name.from_text("hmac-sha1")
-HMAC_SHA224 = dns.name.from_text("hmac-sha224")
-HMAC_SHA256 = dns.name.from_text("hmac-sha256")
-HMAC_SHA384 = dns.name.from_text("hmac-sha384")
-HMAC_SHA512 = dns.name.from_text("hmac-sha512")
-GSS_TSIG = dns.name.from_text("gss-tsig")
-
-_hashes = {
-    HMAC_SHA224: hashlib.sha224,
-    HMAC_SHA256: hashlib.sha256,
-    HMAC_SHA384: hashlib.sha384,
-    HMAC_SHA512: hashlib.sha512,
-    HMAC_SHA1: hashlib.sha1,
-    HMAC_MD5: hashlib.md5,
-}
-
-default_algorithm = HMAC_SHA256
-
-
 def _digest(wire, key, rdata, time=None, request_mac=None, ctx=None,
             multi=None):
     """Return a context containing the TSIG rdata for the input parameters
diff --git a/tests/test_tsig.py b/tests/test_tsig.py
index ec5a6cc..6e3993b 100644
--- a/tests/test_tsig.py
+++ b/tests/test_tsig.py
@@ -226,3 +226,22 @@ class TSIGTestCase(unittest.TestCase):
             def bad():
                 dns.message.from_wire(w, keyring=keyring, request_mac=q.mac)
             self.assertRaises(ex, bad)
+
+    def _test_truncated_algorithm(self, alg, length):
+        key = dns.tsig.Key('foo', b'abcdefg', algorithm=alg)
+        q = dns.message.make_query('example', 'a')
+        q.use_tsig(key)
+        q2 = dns.message.from_wire(q.to_wire(), keyring=key)
+
+        self.assertTrue(q2.had_tsig)
+        self.assertEqual(q2.tsig[0].algorithm, q.tsig[0].algorithm)
+        self.assertEqual(len(q2.tsig[0].mac), length // 8)
+
+    def test_hmac_sha256_128(self):
+        self._test_truncated_algorithm(dns.tsig.HMAC_SHA256_128, 128)
+
+    def test_hmac_sha384_192(self):
+        self._test_truncated_algorithm(dns.tsig.HMAC_SHA384_192, 192)
+
+    def test_hmac_sha512_256(self):
+        self._test_truncated_algorithm(dns.tsig.HMAC_SHA512_256, 256)