summaryrefslogtreecommitdiff
path: root/project2/proj2_s4498062/dns/message.py
diff options
context:
space:
mode:
Diffstat (limited to 'project2/proj2_s4498062/dns/message.py')
-rw-r--r--project2/proj2_s4498062/dns/message.py325
1 files changed, 0 insertions, 325 deletions
diff --git a/project2/proj2_s4498062/dns/message.py b/project2/proj2_s4498062/dns/message.py
deleted file mode 100644
index a8ad621..0000000
--- a/project2/proj2_s4498062/dns/message.py
+++ /dev/null
@@ -1,325 +0,0 @@
-""" DNS messages
-
-This module contains classes for DNS messages, their header section and
-question fields. See section 4 of RFC 1035 for more info.
-"""
-
-import struct
-
-from dns.domainname import Parser, Composer
-from dns.resource import ResourceRecord
-from dns.types import Type
-
-
-class Message(object):
- """ DNS message """
-
- def __init__(
- self, header, questions=None, answers=None, authorities=None,
- additionals=None):
- """ Create a new DNS message
-
- Args:
- header (Header): the header section
- questions ([Question]): the question section
- answers ([ResourceRecord]): the answer section
- authorities ([ResourceRecord]): the authority section
- additionals ([ResourceRecord]): the additional section
- """
- self.header = header
- self.questions = questions if questions != None else []
- self.answers = answers if answers != None else []
- self.authorities = authorities if authorities != None else []
- self.additionals = additionals if additionals != None else []
-
- @property
- def resources(self):
- """ Getter for all resource records """
- return self.answers + self.authorities + self.additionals
-
- def to_bytes(self):
- """ Convert Message to bytes """
- composer = Composer()
-
- # Add header
- result = self.header.to_bytes()
-
- # Add questions
- for question in self.questions:
- offset = len(result)
- result += question.to_bytes(offset, composer)
-
- # Add answers
- for answer in self.answers:
- offset = len(result)
- result += answer.to_bytes(offset, composer)
-
- # Add authorities
- for authority in self.authorities:
- offset = len(result)
- result += authority.to_bytes(offset, composer)
-
- # Add additionals
- for additional in self.additionals:
- offset = len(result)
- result += additional.to_bytes(offset, composer)
-
- return result
-
- @classmethod
- def from_bytes(cls, packet):
- """ Create Message from bytes
-
- Args:
- packet (bytes): byte representation of the message
- """
- parser = Parser()
-
- # Parse header
- header, offset = Header.from_bytes(packet), 12
-
- # Parse questions
- questions = []
- for _ in range(header.qd_count):
- question, offset = Question.from_bytes(packet, offset, parser)
- questions.append(question)
-
- # Parse answers
- answers = []
- for _ in range(header.an_count):
- answer, offset = ResourceRecord.from_bytes(packet, offset, parser)
- answers.append(answer)
-
- # Parse authorities
- authorities = []
- for _ in range(header.ns_count):
- auth, offset = ResourceRecord.from_bytes(packet, offset, parser)
- authorities.append(auth)
-
- # Parse additionals
- additionals = []
- for _ in range(header.ar_count):
- add, offset = ResourceRecord.from_bytes(packet, offset, parser)
- additionals.append(add)
-
- return cls(header, questions, answers, authorities, additionals)
-
- def get_addresses(self):
- """Get the addresses from a response"""
- return self.get_data('answers', Type.A)
-
- def get_hints(self):
- """Get the nameservers from a response"""
- hints = self.get_data('authorities', Type.NS)
- results = []
- for hint in hints:
- addits = [
- a.rdata.data for a in self.additionals if
- a.name == hint and a.type_ == Type.A]
- if addits == []:
- continue
- results.append((hint, addits))
- return results
-
- def get_aliases(self):
- """Get the aliases from a response"""
- return self.get_data('additionals', Type.CNAME)
-
- def get_data(self, key, type_=None):
- """Yield from a message section, possibly filtering by type"""
- for k in getattr(self, key):
- if type_ == None or k.type_ == type_:
- yield k.rdata.data
-
-
-class Header(object):
- """ The header section of a DNS message
-
- Contains a number of properties which are accessible as normal member
- variables.
-
- See section 4.1.1 of RFC 1035 for their meaning.
- """
-
- def __init__(self, ident, flags, qd_count, an_count, ns_count, ar_count):
- """ Create a new Header object
-
- Args:
- ident (int): identifier
- qd_count (int): number of entries in question section
- an_count (int): number of entries in answer section
- ns_count (int): number of entries in authority section
- ar_count (int): number of entries in additional section
- """
- self.ident = ident
- self._flags = flags
- self.qd_count = qd_count
- self.an_count = an_count
- self.ns_count = ns_count
- self.ar_count = ar_count
-
- def to_bytes(self):
- """ Convert header to bytes """
- return struct.pack(
- "!6H",
- self.ident,
- self._flags,
- self.qd_count,
- self.an_count,
- self.ns_count,
- self.ar_count)
-
- @classmethod
- def from_bytes(cls, packet):
- """ Convert Header from bytes """
- if len(packet) < 12:
- raise ShortHeader
- return cls(*struct.unpack_from("!6H", packet))
-
- @property
- def flags(self):
- """The flags of the header"""
- return self._flags
-
- @flags.setter
- def flags(self, value):
- """Set the flags of the header"""
- if value >= (1 << 16):
- raise ValueError("value too big for flags")
- self._flags = value
-
- @property
- def qr(self):
- """The QR flag"""
- return self._flags & (1 << 15)
-
- @qr.setter
- def qr(self, value):
- """Set the QR flag"""
- if value:
- self._flags |= (1 << 15)
- else:
- self._flags &= ~(1 << 15)
-
- @property
- def opcode(self):
- """The opcode of the header"""
- return (self._flags & (((1 << 4) - 1) << 11)) >> 11
-
- @opcode.setter
- def opcode(self, value):
- """Set the opcode"""
- if value > 0b1111:
- raise ValueError("invalid opcode")
- self._flags &= ~(((1 << 4) - 1) << 11)
- self._flags |= value << 11
-
- @property
- def aa(self):
- """The AA flag"""
- return self._flags & (1 << 10)
-
- @aa.setter
- def aa(self, value):
- """Set the AA flag"""
- if value:
- self._flags |= (1 << 10)
- else:
- self._flags &= ~(1 << 10)
-
- @property
- def tc(self):
- """The TC flag"""
- return self._flags & (1 << 9)
-
- @tc.setter
- def tc(self, value):
- """Set the TC flag"""
- if value:
- self._flags |= (1 << 9)
- else:
- self._flags &= ~(1 << 9)
-
- @property
- def rd(self):
- """The RD flag"""
- return self._flags & (1 << 8)
-
- @rd.setter
- def rd(self, value):
- """Set the RD flag"""
- if value:
- self._flags |= (1 << 8)
- else:
- self._flags &= ~(1 << 8)
-
- @property
- def ra(self):
- """The RA flag"""
- return self._flags & (1 << 7)
-
- @ra.setter
- def ra(self, value):
- """Set the RA flag"""
- if value:
- self._flags |= (1 << 7)
- else:
- self._flags &= ~(1 << 7)
-
- @property
- def z(self):
- """The Z flag"""
- return self._flags & (((1 << 3) - 1) << 4) >> 4
-
- @z.setter
- def z(self, value):
- """Set the Z flag"""
- if value:
- raise ValueError("non-zero zero flag")
-
- @property
- def rcode(self):
- """The return code"""
- return self._flags & ((1 << 4) - 1)
-
- @rcode.setter
- def rcode(self, value):
- """Set the return code"""
- if value > 0b1111:
- raise ValueError("invalid return code")
- self._flags &= ~((1 << 4) - 1)
- self._flags |= value
-
-
-class Question(object):
- """ An entry in the question section.
-
- See section 4.1.2 of RFC 1035 for more info.
- """
-
- def __init__(self, qname, qtype, qclass):
- """ Create a new entry in the question section
-
- Args:
- qname (str): QNAME
- qtype (Type): QTYPE
- qclass (Class): QCLASS
- """
- self.qname = qname
- self.qtype = qtype
- self.qclass = qclass
-
- def to_bytes(self, offset, composer):
- """ Convert Question to bytes """
- bqname = composer.to_bytes(offset, [self.qname])
- bqtype = struct.pack("!H", self.qtype)
- bqclass = struct.pack("!H", self.qclass)
- return bqname + bqtype + bqclass
-
- @classmethod
- def from_bytes(cls, packet, offset, parser):
- """ Convert Question from bytes """
- qnames, offset = parser.from_bytes(packet, offset, 1)
- qname = qnames[0]
- qtype, qclass = struct.unpack_from("!2H", packet, offset)
- return cls(qname, qtype, qclass), offset + 4