tlsfuzzer/tlslite-ng

# Copyright (c) 2015, Hubert Kario
#
# See the LICENSE file for legal information regarding use of this file.

"""Helper package for handling fragmentation of messages."""

from __future__ import generators

from .utils.codec import Parser
from .utils.deprecations import deprecated_attrs, deprecated_params


@deprecated_attrs({"add_static_size": "addStaticSize",
                   "add_dynamic_size": "addDynamicSize",
                   "add_data": "addData",
                   "get_message": "getMessage",
                   "clear_buffers": "clearBuffers"})
class Defragmenter(object):
    """
    Class for demultiplexing TLS messages.

    Since the messages can be interleaved and fragmented between each other
    we need to cache not complete ones and return in order of urgency.

    Supports messages with given size (like Alerts) or with a length header
    in specific place (like Handshake messages).

    :ivar priorities: order in which messages from given types should be
        returned.
    :ivar buffers: data buffers for message types
    :ivar decoders: functions which check buffers if a message of given type
        is complete
    """

    def __init__(self):
        """Set up empty defregmenter"""
        self.priorities = []
        self.buffers = {}
        self.decoders = {}

    @deprecated_params({"msg_type": "msgType"})
    def add_static_size(self, msg_type, size):
        """Add a message type which all messages are of same length"""
        if msg_type in self.priorities:
            raise ValueError("Message type already defined")
        if size < 1:
            raise ValueError("Message size must be positive integer")

        self.priorities += [msg_type]

        self.buffers[msg_type] = bytearray(0)
        def size_handler(data):
            """
            Size of message in parameter

            If complete message is present in parameter returns its size,
            None otherwise.
            """
            if len(data) < size:
                return None
            else:
                return size
        self.decoders[msg_type] = size_handler

    @deprecated_params({"msg_type": "msgType",
                        "size_offset": "sizeOffset",
                        "size_of_size": "sizeOfSize"})
    def add_dynamic_size(self, msg_type, size_offset, size_of_size):
        """Add a message type which has a dynamic size set in a header"""
        if msg_type in self.priorities:
            raise ValueError("Message type already defined")
        if size_of_size < 1:
            raise ValueError("Size of size must be positive integer")
        if size_offset < 0:
            raise ValueError("Offset can't be negative")

        self.priorities += [msg_type]
        self.buffers[msg_type] = bytearray(0)

        def size_handler(data):
            """
            Size of message in parameter

            If complete message is present in parameter returns its size,
            None otherwise.
            """
            if len(data) < size_offset+size_of_size:
                return None
            else:
                parser = Parser(data)
                # skip the header
                parser.skip_bytes(size_offset)

                payload_length = parser.get(size_of_size)
                if parser.getRemainingLength() < payload_length:
                    # not enough bytes in buffer
                    return None
                return size_offset + size_of_size + payload_length

        self.decoders[msg_type] = size_handler

    @deprecated_params({"msg_type": "msgType"})
    def add_data(self, msg_type, data):
        """Adds data to buffers"""
        if msg_type not in self.priorities:
            raise ValueError("Message type not defined")

        self.buffers[msg_type] += data

    def get_message(self):
        """Extract the highest priority complete message from buffer"""
        for msg_type in self.priorities:
            buf = self.buffers[msg_type]
            length = self.decoders[msg_type](buf)
            if length is None:
                continue

            # extract message
            data = buf[:length]
            # remove it from buffer
            del buf[:length]
            return (msg_type, data)
        return None

    def clear_buffers(self):
        """Remove all data from buffers"""
        for key in self.buffers.keys():
            self.buffers[key] = bytearray(0)

    def is_empty(self):
        """Return True if all buffers are empty."""
        return all(not i for i in self.buffers.values())