raise an exception if the TSIG error is non-zero

1 files changed, 32 insertions, 3 deletions

diff --git a/dns/tsig.py b/dns/tsig.py
index 247e7d3..5c112ac 100644
--- a/dns/tsig.py
+++ b/dns/tsig.py
@@ -30,8 +30,28 @@ class BadSignature(dns.exception.DNSException):
     """Raised if the TSIG signature fails to verify."""
     pass
 
+class PeerError(dns.exception.DNSException):
+    """Base class for all TSIG errors generated by the remote peer"""
+    pass
+
+class PeerBadKey(PeerError):
+    """Raised if the peer didn't know the key we used"""
+    pass
+
+class PeerBadSignature(PeerError):
+    """Raised if the peer didn't like the signature we sent"""
+    pass
+
+class PeerBadTime(PeerError):
+    """Raised if the peer didn't like the time we sent"""
+    pass
+
 _alg_name = dns.name.from_text('HMAC-MD5.SIG-ALG.REG.INT.').to_digestable()
-    
+
+BADSIG = 16
+BADKEY = 17
+BADTIME = 18
+
 def hmac_md5(wire, keyname, secret, time, fudge, original_id, error,
              other_data, request_mac, ctx=None, multi=False, first=True):
     """Return a (tsig_rdata, mac, ctx) tuple containing the HMAC-MD5 TSIG rdata
@@ -40,7 +60,7 @@ def hmac_md5(wire, keyname, secret, time, fudge, original_id, error,
     @rtype: (string, string, hmac.HMAC object)
     @raises ValueError: I{other_data} is too long
     """
-    
+
     if first:
         ctx = hmac.new(secret)
         ml = len(request_mac)
@@ -57,7 +77,7 @@ def hmac_md5(wire, keyname, secret, time, fudge, original_id, error,
     long_time = time + 0L
     upper_time = (long_time >> 32) & 0xffffL
     lower_time = long_time & 0xffffffffL
-    time_mac = struct.pack('!HIH', upper_time, lower_time, fudge) 
+    time_mac = struct.pack('!HIH', upper_time, lower_time, fudge)
     pre_mac = _alg_name + time_mac
     ol = len(other_data)
     if ol > 65535:
@@ -111,6 +131,15 @@ def validate(wire, keyname, secret, now, request_mac, tsig_start, tsig_rdata,
     current += other_size
     if current != tsig_rdata + tsig_rdlen:
         raise dns.exception.FormError
+    if error != 0:
+        if error == BADSIG:
+            raise PeerBadSignature
+        elif error == BADKEY:
+            raise PeerBadKey
+        elif error == BADTIME:
+            raise PeerBadTime
+        else:
+            raise PeerError, 'unknown TSIG error code %d' % error
     time_low = time - fudge
     time_high = time + fudge
     if now < time_low or now > time_high: