Merge pull request #71 from spacekpe/DNSKEY.flags_to_text_set

Add dns.rdtypes.ANY.DNSKEY.flags_to_text_set()

2 files changed, 116 insertions, 0 deletions

diff --git a/dns/rdtypes/ANY/DNSKEY.py b/dns/rdtypes/ANY/DNSKEY.py
index 1d678d2..7bc5850 100644
--- a/dns/rdtypes/ANY/DNSKEY.py
+++ b/dns/rdtypes/ANY/DNSKEY.py
@@ -20,11 +20,54 @@ import dns.exception
 import dns.dnssec
 import dns.rdata
 
+
 # flag constants
 SEP = 0x0001
 REVOKE = 0x0080
 ZONE = 0x0100
 
+_flag_by_text = {
+    'SEP': SEP,
+    'REVOKE': REVOKE,
+    'ZONE': ZONE
+    }
+
+# We construct the inverse mapping programmatically to ensure that we
+# cannot make any mistakes (e.g. omissions, cut-and-paste errors) that
+# would cause the mapping not to be true inverse.
+_flag_by_value = dict([(y, x) for x, y in _flag_by_text.iteritems()])
+
+
+def flags_to_text_set(flags):
+    """Convert a DNSKEY flags value to set texts
+    @rtype: set([string])"""
+
+    flags_set = set()
+    mask = 0x1
+    while mask <= 0x8000:
+        if flags & mask:
+            text = _flag_by_value.get(mask)
+            if not text:
+                text = hex(mask)
+            flags_set.add(text)
+        mask <<= 1
+    return flags_set
+
+
+def flags_from_text_set(texts_set):
+    """Convert set of DNSKEY flag mnemonic texts to DNSKEY flag value
+    @rtype: int"""
+
+    flags = 0
+    for text in texts_set:
+        try:
+            flags += _flag_by_text[text]
+        except KeyError:
+            raise NotImplementedError(
+                "DNSKEY flag '%s' is not supported" % text)
+    return flags
+
+
 class DNSKEY(dns.rdata.Rdata):
     """DNSKEY record
 
@@ -92,3 +135,8 @@ class DNSKEY(dns.rdata.Rdata):
         if v == 0:
             v = cmp(self.key, other.key)
         return v
+
+    def flags_to_text_set(self):
+        """Convert a DNSKEY flags value to set texts
+        @rtype: set([string])"""
+        return flags_to_text_set(self.flags)
diff --git a/tests/test_rdtypeanydnskey.py b/tests/test_rdtypeanydnskey.py
new file mode 100644
index 0000000..d9e40d7
--- /dev/null
+++ b/tests/test_rdtypeanydnskey.py
@@ -0,0 +1,68 @@
+# Copyright (C) 2014 Red Hat, Inc.
+# Author: Petr Spacek <pspacek@redhat.com>
+#
+# Permission to use, copy, modify, and distribute this software and its
+# documentation for any purpose with or without fee is hereby granted,
+# provided that the above copyright notice and this permission notice
+# appear in all copies.
+#
+# THE SOFTWARE IS PROVIDED 'AS IS' AND RED HAT DISCLAIMS ALL WARRANTIES
+# WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
+# MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL NOMINUM BE LIABLE FOR
+# ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
+# WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
+# ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT
+# OF OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
+
+import unittest
+
+import dns.rrset
+import dns.rdtypes.ANY.DNSKEY
+
+
+class RdtypeAnyDnskeyTestCase(unittest.TestCase):
+
+    def testFlagsEmpty(self):
+        '''Test DNSKEY flag to/from text conversion for zero flag/empty set.'''
+        good_s = set()
+        good_f = 0
+        from_flags = dns.rdtypes.ANY.DNSKEY.flags_to_text_set(good_f)
+        self.failUnless(from_flags == good_s,
+                        '"%s" != "%s"' % (from_flags, good_s))
+        from_set = dns.rdtypes.ANY.DNSKEY.flags_from_text_set(good_s)
+        self.failUnless(from_set == good_f,
+                        '"0x%x" != "0x%x"' % (from_set, good_f))
+
+    def testFlagsAll(self):
+        '''Test that all defined flags are recognized.'''
+        good_s = set(['SEP', 'REVOKE', 'ZONE'])
+        good_f = 0x181
+        from_flags = dns.rdtypes.ANY.DNSKEY.flags_to_text_set(good_f)
+        self.failUnless(from_flags == good_s,
+                        '"%s" != "%s"' % (from_flags, good_s))
+        from_text = dns.rdtypes.ANY.DNSKEY.flags_from_text_set(good_s)
+        self.failUnless(from_text == good_f,
+                        '"0x%x" != "0x%x"' % (from_text, good_f))
+
+    def testFlagsUnknownToText(self):
+        '''Test that undefined flags are returned in hexadecimal notation.'''
+        unk_s = set(['0x8000'])
+        flags_s = dns.rdtypes.ANY.DNSKEY.flags_to_text_set(0x8000)
+        self.failUnless(flags_s == unk_s, '"%s" != "%s"' % (flags_s, unk_s))
+
+    def testFlagsUnknownToFlags(self):
+        '''Test that conversion from undefined mnemonic raises error.'''
+        self.failUnlessRaises(NotImplementedError,
+                              dns.rdtypes.ANY.DNSKEY.flags_from_text_set,
+                              (['0x8000']))
+
+    def testFlagsRRToText(self):
+        '''Test that RR method returns correct flags.'''
+        rr = dns.rrset.from_text('foo', 300, 'IN', 'DNSKEY', '257 3 8 KEY=')[0]
+        rr_s = set(['ZONE', 'SEP'])
+        flags_s = rr.flags_to_text_set()
+        self.failUnless(flags_s == rr_s, '"%s" != "%s"' % (flags_s, rr_s))
+
+
+if __name__ == '__main__':
+    unittest.main()