processing order implementation

9 files changed, 133 insertions, 0 deletions

diff --git a/dns/rdata.py b/dns/rdata.py
index acf34ae..7846205 100644
--- a/dns/rdata.py
+++ b/dns/rdata.py
@@ -23,6 +23,7 @@ import binascii
 import io
 import inspect
 import itertools
+import random
 
 import dns.wire
 import dns.exception
@@ -459,6 +460,14 @@ class Rdata:
             # against *as_value*.
             return tuple(as_value(v) for v in value)
 
+    # Processing order
+
+    @classmethod
+    def _processing_order(cls, iterable):
+        items = list(iterable)
+        random.shuffle(items)
+        return items
+
 
 class GenericRdata(Rdata):
 
diff --git a/dns/rdataset.py b/dns/rdataset.py
index 2e3b4d4..f27e251 100644
--- a/dns/rdataset.py
+++ b/dns/rdataset.py
@@ -306,6 +306,20 @@ class Rdataset(dns.set.Set):
             return True
         return False
 
+    def processing_order(self):
+        """Return rdatas in a valid processing order according to the type's
+        specification.  For example, MX records are in preference order from
+        lowest to highest preferences, with items of the same perference
+        shuffled.
+
+        For types that do not define a processing order, the rdatas are
+        simply shuffled.
+        """
+        if len(self) == 0:
+            return []
+        else:
+            return self[0]._processing_order(iter(self))
+
 
 @dns.immutable.immutable
 class ImmutableRdataset(Rdataset):
diff --git a/dns/rdtypes/ANY/URI.py b/dns/rdtypes/ANY/URI.py
index 60a43c8..4e7ad01 100644
--- a/dns/rdtypes/ANY/URI.py
+++ b/dns/rdtypes/ANY/URI.py
@@ -67,3 +67,13 @@ class URI(dns.rdata.Rdata):
         if len(target) == 0:
             raise dns.exception.FormError('URI target may not be empty')
         return cls(rdclass, rdtype, priority, weight, target)
+
+    def _processing_priority(self):
+        return self.priority
+
+    def _processing_weight(self):
+        return self.weight
+
+    @classmethod
+    def _processing_order(cls, iterable):
+        return dns.rdtypes.util.weighted_processing_order(iterable, False)
diff --git a/dns/rdtypes/IN/NAPTR.py b/dns/rdtypes/IN/NAPTR.py
index a4058ad..1f072f8 100644
--- a/dns/rdtypes/IN/NAPTR.py
+++ b/dns/rdtypes/IN/NAPTR.py
@@ -89,3 +89,10 @@ class NAPTR(dns.rdata.Rdata):
         replacement = parser.get_name(origin)
         return cls(rdclass, rdtype, order, preference, strings[0], strings[1],
                    strings[2], replacement)
+
+    def _processing_priority(self):
+        return (self.order, self.preference)
+
+    @classmethod
+    def _processing_order(cls, iterable):
+        return dns.rdtypes.util.priority_processing_order(iterable)
diff --git a/dns/rdtypes/IN/PX.py b/dns/rdtypes/IN/PX.py
index 3a744f6..8abfb29 100644
--- a/dns/rdtypes/IN/PX.py
+++ b/dns/rdtypes/IN/PX.py
@@ -63,3 +63,10 @@ class PX(dns.rdata.Rdata):
         map822 = parser.get_name(origin)
         mapx400 = parser.get_name(origin)
         return cls(rdclass, rdtype, preference, map822, mapx400)
+
+    def _processing_priority(self):
+        return self.preference
+
+    @classmethod
+    def _processing_order(cls, iterable):
+        return dns.rdtypes.util.priority_processing_order(iterable)
diff --git a/dns/rdtypes/IN/SRV.py b/dns/rdtypes/IN/SRV.py
index 99ff70b..3eb227f 100644
--- a/dns/rdtypes/IN/SRV.py
+++ b/dns/rdtypes/IN/SRV.py
@@ -63,3 +63,13 @@ class SRV(dns.rdata.Rdata):
         (priority, weight, port) = parser.get_struct('!HHH')
         target = parser.get_name(origin)
         return cls(rdclass, rdtype, priority, weight, port, target)
+
+    def _processing_priority(self):
+        return self.priority
+
+    def _processing_weight(self):
+        return self.weight
+
+    @classmethod
+    def _processing_order(cls, iterable):
+        return dns.rdtypes.util.weighted_processing_order(iterable, True)
diff --git a/dns/rdtypes/mxbase.py b/dns/rdtypes/mxbase.py
index 7bacbd8..5641823 100644
--- a/dns/rdtypes/mxbase.py
+++ b/dns/rdtypes/mxbase.py
@@ -23,6 +23,7 @@ import dns.exception
 import dns.immutable
 import dns.rdata
 import dns.name
+import dns.rdtypes.util
 
 
 @dns.immutable.immutable
@@ -59,6 +60,13 @@ class MXBase(dns.rdata.Rdata):
         exchange = parser.get_name(origin)
         return cls(rdclass, rdtype, preference, exchange)
 
+    def _processing_priority(self):
+        return self.preference
+
+    @classmethod
+    def _processing_order(cls, iterable):
+        return dns.rdtypes.util.priority_processing_order(iterable)
+
 
 @dns.immutable.immutable
 class UncompressedMX(MXBase):
diff --git a/dns/rdtypes/svcbbase.py b/dns/rdtypes/svcbbase.py
index 4a3a217..585c05e 100644
--- a/dns/rdtypes/svcbbase.py
+++ b/dns/rdtypes/svcbbase.py
@@ -534,3 +534,10 @@ class SVCBBase(dns.rdata.Rdata):
                 value = pcls.from_wire_parser(parser, origin)
             params[key] = value
         return cls(rdclass, rdtype, priority, target, params)
+
+    def _processing_priority(self):
+        return self.priority
+
+    @classmethod
+    def _processing_order(cls, iterable):
+        return dns.rdtypes.util.priority_processing_order(iterable)
diff --git a/dns/rdtypes/util.py b/dns/rdtypes/util.py
index 30be37d..c2fb603 100644
--- a/dns/rdtypes/util.py
+++ b/dns/rdtypes/util.py
@@ -15,6 +15,7 @@
 # ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT
 # OF OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 
+import random
 import struct
 
 import dns.exception
@@ -183,3 +184,63 @@ class Bitmap:
             bitmap = parser.get_counted_bytes()
             windows.append((window, bitmap))
         return cls(windows)
+
+
+def _priority_table(items):
+    by_priority = {}
+    for rdata in items:
+        key = rdata._processing_priority()
+        rdatas = by_priority.get(key)
+        if rdatas is None:
+            rdatas = []
+            by_priority[key] = rdatas
+        rdatas.append(rdata)
+    return by_priority
+
+def priority_processing_order(iterable):
+    items = list(iterable)
+    if len(items) == 1:
+        return [items[0]]
+    by_priority = _priority_table(items)
+    ordered = []
+    for k in sorted(by_priority.keys()):
+        rdatas = by_priority[k]
+        random.shuffle(rdatas)
+        ordered.extend(rdatas)
+    return ordered
+
+def _processing_weight(rdata, adjust_zero_weight):
+    weight = rdata._processing_weight()
+    if weight == 0 and adjust_zero_weight:
+        return 0.1
+    else:
+        return weight
+
+def weighted_processing_order(iterable, adjust_zero_weight=False):
+    items = list(iterable)
+    if len(items) == 1:
+        return [items[0]]
+    by_priority = _priority_table(items)
+    ordered = []
+    for k in sorted(by_priority.keys()):
+        weights_vary = False
+        weights = []
+        rdatas = by_priority[k]
+        for rdata in rdatas:
+            weight = _processing_weight(rdata, adjust_zero_weight)
+            if len(weights) > 0 and weight != weights[-1]:
+                weights_vary = True
+            weights.append(weight)
+        if weights_vary:
+            while len(rdatas) > 1:
+                items = random.choices(rdatas, weights)
+                rdata = items[0]
+                ordered.append(rdata)
+                rdatas.remove(rdata)
+                weight = _processing_weight(rdata, adjust_zero_weight)
+                weights.remove(weight)
+            ordered.append(rdatas[0])
+        else:
+            random.shuffle(rdatas)
+            ordered.extend(rdatas)
+    return ordered