Composer.py 2.82 KB
"""
 @author Georg Hopp

"""
import struct

class Composer(object):
    def __init__(self):
        self._name_ofs = {}

    def compose(self, message):
        self._name_ofs = {}

        header = struct.pack(
            '!HHHHHH',
            message._msg_id,
            message._flags,
            len(message._queries),
            len(message._answers),
            len(message._authoritys),
            len(message._additionals)
        )

        queries = answers = authoritys = additionals = ''
        ofs = len(header)
        if message._queries:
            queries = self._composeQueries(message, ofs)

        ofs += len(queries)
        if message._answers:
            answers = self._composeAnswers(message, ofs)

        ofs += len(answers)
        if message._authoritys:
            authoritys = self._composeAuthoritys(message, ofs)

        ofs += len(authoritys)
        if message._additionals:
            additionals = self._composeAdditionals(message, ofs)

        return header + queries + answers + authoritys + additionals

    def _composeQueries(self, message, ofs):
        encoded = ''

        for query in message._queries:
            name, typ, cls = query
            ename = self._encodeName(name, ofs)

            query    = struct.pack('!%dsHH'%len(ename), ename, typ, cls)
            ofs     += len(query)
            encoded += query

        return encoded

    def _composeAnswers(self, message, ofs):
        encoded = ''

        for answer in message._answers:
            record   = self._composeResourceRecord(answer, ofs)
            ofs     += len(record)
            encoded += record

        return encoded

    def _composeAuthoritys(self, message, ofs):
        encoded = ''

        for authority in message._authoritys:
            record   = self._composeResourceRecord(authority, ofs)
            ofs     += len(record)
            encoded += record

        return encoded

    def _composeAdditionals(self, message, ofs):
        encoded = ''

        for additional in message._additionals:
            record   = self._composeResourceRecord(additional, ofs)
            ofs     += len(record)
            encoded += record

        return encoded

    def _composeResourceRecord(self, record, ofs):
        name, typ, cls, ttl, data = record
        ename = self._encodeName(name, ofs)
        return struct.pack('!%dsHHLH%ds'%(len(ename), len(data)),
                ename, typ, cls, ttl, len(data), data)

    def _encodeName(self, name, ofs):
        if name in self._name_ofs:
            name = struct.pack('!H',
                int('1100000000000000', 2) | self._name_ofs[name])
        else:
            self._name_ofs[name] = ofs
            name = ''.join([struct.pack('B%ds'%len(p), len(p), p)
                for p in name.split('.')]) + '\x00'

        return name

# vim: set ft=python et ts=4 sw=4 sts=4: