From 005e472d864d84838fff6e96a16f63f9a2a7a5f5 Mon Sep 17 00:00:00 2001 From: Bob Halley Date: Thu, 25 Jun 2020 17:13:05 -0700 Subject: [PATCH] make sections a list; propertize question, answer, etc. --- dns/message.py | 108 +++++++++++++++++++++++++++---------------------- dns/update.py | 50 ++++++++++++++++++----- 2 files changed, 98 insertions(+), 60 deletions(-) diff --git a/dns/message.py b/dns/message.py index 638f2048..eabca4a0 100644 --- a/dns/message.py +++ b/dns/message.py @@ -85,6 +85,11 @@ class MessageSection(dns.enum.IntEnum): AUTHORITY = 2 ADDITIONAL = 3 + @classmethod + def _maximum(cls): + return 3 + + globals().update(MessageSection.__members__) class Message: @@ -98,10 +103,7 @@ class Message: else: self.id = id self.flags = 0 - self.question = [] - self.answer = [] - self.authority = [] - self.additional = [] + self.sections = [[], [], [], []] self.edns = -1 self.ednsflags = 0 self.payload = 0 @@ -124,6 +126,38 @@ class Message: self.first = True self.index = {} + @property + def question(self): + return self.sections[0] + + @question.setter + def question(self, v): + self.sections[0] = v + + @property + def answer(self): + return self.sections[1] + + @answer.setter + def answer(self, v): + self.sections[1] = v + + @property + def authority(self): + return self.sections[2] + + @authority.setter + def authority(self, v): + self.sections[2] = v + + @property + def additional(self): + return self.sections[3] + + @additional.setter + def additional(self, v): + self.sections[3] = v + def __repr__(self): return '' @@ -182,24 +216,15 @@ class Message: return False if self.flags != other.flags: return False - for n in self.question: - if n not in other.question: - return False - for n in other.question: - if n not in self.question: - return False - for n in self.answer: - if n not in other.answer: - return False - for n in other.answer: - if n not in self.answer: - return False - for n in self.authority: - if n not in other.authority: - return False - for n in other.authority: - if n not in self.authority: - return False + for i in range(4): + section = self.sections[i] + other_section = other.sections[i] + for n in section: + if n not in other_section: + return False + for n in other_section: + if n not in section: + return False return True def __ne__(self, other): @@ -235,8 +260,7 @@ class Message: def section_number(self, section): """Return the "section number" of the specified section for use - in indexing. The question section is 0, the answer section is 1, - the authority section is 2, and the additional section is 3. + in indexing. *section* is one of the section attributes of this message. @@ -245,39 +269,25 @@ class Message: Returns an ``int``. """ - if section is self.question: - return QUESTION - elif section is self.answer: - return ANSWER - elif section is self.authority: - return AUTHORITY - elif section is self.additional: - return ADDITIONAL - else: - raise ValueError('unknown section') + for i in range(4): + if section is self.sections[i]: + return self._section_enum(i) + raise ValueError('unknown section') def section_from_number(self, number): - """Return the "section number" of the specified section for use - in indexing. The question section is 0, the answer section is 1, - the authority section is 2, and the additional section is 3. + """Return the section list associated with the specified section + number. - *section* is one of the section attributes of this message. + *number* is a section number `int` or the text form of a section + name. Raises ``ValueError`` if the section isn't known. - Returns an ``int``. + Returns a ``list``. """ - if number == QUESTION: - return self.question - elif number == ANSWER: - return self.answer - elif number == AUTHORITY: - return self.authority - elif number == ADDITIONAL: - return self.additional - else: - raise ValueError('unknown section') + section = self._section_enum.make(number) + return self.sections[section] def find_rrset(self, section, name, rdclass, rdtype, covers=dns.rdatatype.NONE, deleting=None, create=False, diff --git a/dns/update.py b/dns/update.py index 4878db71..05e9a530 100644 --- a/dns/update.py +++ b/dns/update.py @@ -34,6 +34,10 @@ class UpdateSection(dns.enum.IntEnum): UPDATE = 2 ADDITIONAL = 3 + @classmethod + def _maximum(cls): + return 3 + class Update(dns.message.Message): @@ -70,16 +74,40 @@ class Update(dns.message.Message): self.origin = zone rdclass = dns.rdataclass.RdataClass.make(rdclass) self.zone_rdclass = rdclass - self.find_rrset(self.question, self.origin, rdclass, dns.rdatatype.SOA, + self.find_rrset(self.zone, self.origin, rdclass, dns.rdatatype.SOA, create=True, force_unique=True) if keyring is not None: self.use_tsig(keyring, keyname, algorithm=keyalgorithm) + @property + def zone(self): + return self.sections[0] + + @zone.setter + def zone(self, v): + self.sections[0] = v + + @property + def prerequisite(self): + return self.sections[1] + + @prerequisite.setter + def prerequisite(self, v): + self.sections[1] = v + + @property + def update(self): + return self.sections[2] + + @update.setter + def update(self, v): + self.sections[2] = v + def _add_rr(self, name, ttl, rd, deleting=None, section=None): """Add a single RR to the update section.""" if section is None: - section = self.authority + section = self.update covers = rd.covers() rrset = self.find_rrset(section, name, self.zone_rdclass, rd.rdtype, covers, deleting, True, True) @@ -139,7 +167,7 @@ class Update(dns.message.Message): - ttl, rdtype, string... """ - self._add(False, self.authority, name, *args) + self._add(False, self.update, name, *args) def delete(self, name, *args): """Delete records. @@ -159,7 +187,7 @@ class Update(dns.message.Message): if isinstance(name, str): name = dns.name.from_text(name, None) if len(args) == 0: - self.find_rrset(self.authority, name, dns.rdataclass.ANY, + self.find_rrset(self.update, name, dns.rdataclass.ANY, dns.rdatatype.ANY, dns.rdatatype.NONE, dns.rdatatype.ANY, True, True) elif isinstance(args[0], dns.rdataset.Rdataset): @@ -174,7 +202,7 @@ class Update(dns.message.Message): else: rdtype = dns.rdatatype.RdataType.make(args.pop(0)) if len(args) == 0: - self.find_rrset(self.authority, name, + self.find_rrset(self.update, name, self.zone_rdclass, rdtype, dns.rdatatype.NONE, dns.rdataclass.ANY, @@ -201,7 +229,7 @@ class Update(dns.message.Message): a delete of the name followed by one or more calls to add. """ - self._add(True, self.authority, name, *args) + self._add(True, self.update, name, *args) def present(self, name, *args): """Require that an owner name (and optionally an rdata type, @@ -221,7 +249,7 @@ class Update(dns.message.Message): if isinstance(name, str): name = dns.name.from_text(name, None) if len(args) == 0: - self.find_rrset(self.answer, name, + self.find_rrset(self.prerequisite, name, dns.rdataclass.ANY, dns.rdatatype.ANY, dns.rdatatype.NONE, None, True, True) @@ -232,10 +260,10 @@ class Update(dns.message.Message): # Add a 0 TTL args = list(args) args.insert(0, 0) - self._add(False, self.answer, name, *args) + self._add(False, self.prerequisite, name, *args) else: rdtype = dns.rdatatype.RdataType.make(args[0]) - self.find_rrset(self.answer, name, + self.find_rrset(self.prerequisite, name, dns.rdataclass.ANY, rdtype, dns.rdatatype.NONE, None, True, True) @@ -247,13 +275,13 @@ class Update(dns.message.Message): if isinstance(name, str): name = dns.name.from_text(name, None) if rdtype is None: - self.find_rrset(self.answer, name, + self.find_rrset(self.prerequisite, name, dns.rdataclass.NONE, dns.rdatatype.ANY, dns.rdatatype.NONE, None, True, True) else: rdtype = dns.rdatatype.RdataType.make(rdtype) - self.find_rrset(self.answer, name, + self.find_rrset(self.prerequisite, name, dns.rdataclass.NONE, rdtype, dns.rdatatype.NONE, None, True, True) -- 2.47.3