Support callable() TSIG keyrings for use-cases like GSSTSig.

1 files changed, 9 insertions, 3 deletions

diff --git a/dns/message.py b/dns/message.py
index 152fa50..4b4eb4d 100644
--- a/dns/message.py
+++ b/dns/message.py
@@ -488,8 +488,8 @@ class Message:
         *key*, a ``dns.tsig.Key`` is the key to use.  If a key is specified,
         the *keyring* and *algorithm* fields are not used.
 
-        *keyring*, a ``dict`` or ``dns.tsig.Key``, is either the TSIG
-        keyring or key to use.
+        *keyring*, a ``dict``, ``callable`` or ``dns.tsig.Key``, is either
+        the TSIG keyring or key to use.
 
         The format of a keyring dict is a mapping from TSIG key name, as
         ``dns.name.Name`` to ``dns.tsig.Key`` or a TSIG secret, a ``bytes``.
@@ -497,7 +497,9 @@ class Message:
         used will be the first key in the *keyring*.  Note that the order of
         keys in a dictionary is not defined, so applications should supply a
         keyname when a ``dict`` keyring is used, unless they know the keyring
-        contains only one key.
+        contains only one key.  If a ``callable`` keyring is specified, the
+        callable will be called with the message and the keyname, and is
+        expected to return a key.
 
         *keyname*, a ``dns.name.Name``, ``str`` or ``None``, the name of
         thes TSIG key to use; defaults to ``None``.  If *keyring* is a
@@ -519,6 +521,8 @@ class Message:
 
         if isinstance(keyring, dns.tsig.Key):
             self.keyring = keyring
+        elif callable(keyring):
+            self.keyring = keyring(self, keyname)
         else:
             if isinstance(keyname, str):
                 keyname = dns.name.from_text(keyname)
@@ -920,6 +924,8 @@ class _WireReader:
                     key = self.keyring.get(absolute_name)
                     if isinstance(key, bytes):
                         key = dns.tsig.Key(absolute_name, key, rd.algorithm)
+                elif callable(self.keyring):
+                    key = self.keyring(self.message, absolute_name)
                 else:
                     key = self.keyring
                 if key is None: