tlsfuzzer/tlslite-ng

# Copyright 2021 Sidney Markowitz All Rights Reserved.
# Distributed under MIT license.
# See file LICENSE for detail or copy at https://opensource.org/licenses/MIT


class BrotliBitReader:
    """
    Wrap a bytes buffer to enable reading 0 < n <=24 bits at a time, or
    transfer of arbitrary number of bytes
    """

    kBitMask = [
        0x000000, 0x000001, 0x000003, 0x000007, 0x00000f, 0x00001f, 0x00003f,
        0x00007f, 0x0000ff, 0x0001ff, 0x0003ff, 0x0007ff, 0x000fff, 0x001fff,
        0x003fff, 0x007fff, 0x00ffff, 0x01ffff, 0x03ffff, 0x07ffff, 0x0fffff,
        0x1fffff, 0x3fffff, 0x7fffff, 0xffffff
    ]

    def __init__(self, input_buffer):
        self.buf_ = bytearray(input_buffer)
        self.buf_len_ = len(input_buffer)
        self.pos_ = 0          # byte position in stream
        # current bit-reading position in current byte (number bits already
        # read from byte, 0-7)
        self.bit_pos_ = 0

    def reset(self):
        """Reset an initialized BrotliBitReader to start of input buffer"""
        self.pos_ = 0
        self.bit_pos_ = 0

    def read_bits(self, n_bits, bits_to_skip=None):
        """
        Get n_bits unsigned integer treating input as little-endian byte
        stream, maybe advancing input buffer pointer

        n_bits: is number of bits to read from input buffer. Set to None or 0
            to seek ahead ignoring the value
        bits_to_skip: number of bits to advance in input_buffer, defaults to
            n_bits if it is None pass in 0 to peek at the next n_bits of value
            without advancing

        It is ok to have n_bits and bits_to_skip be different non-zero values
        if that is what is wanted

        Returns: the next n_bits from the buffer as a little-endian integer,
            0 if n_bits is None or 0
        """
        val = 0
        if bits_to_skip is None:
            bits_to_skip = n_bits
        if n_bits:
            bytes_shift = 0
            buf_pos = self.pos_
            bit_pos_when_done = n_bits + self.bit_pos_
            while bytes_shift < bit_pos_when_done:
                if buf_pos >= self.buf_len_:
                    # if hit end of buffer, this simulates zero padding after
                    # end, which is correct
                    break
                val |= self.buf_[buf_pos] << bytes_shift
                bytes_shift += 8
                buf_pos += 1
            val = (val >> self.bit_pos_) & self.kBitMask[n_bits]
        if bits_to_skip:
            next_in_bits = self.bit_pos_ + bits_to_skip
            self.bit_pos_ = next_in_bits & 7
            self.pos_ += next_in_bits >> 3
        return val

    def copy_bytes(self, dest_buffer, dest_pos, n_bytes):
        """
        Copy bytes from input buffer. This will first skip to next byte
        boundary if not already on one
        """
        if self.bit_pos_ != 0:
            self.bit_pos_ = 0
            self.pos_ += 1
        # call with n_bytes == 0 to just skip to next byte boundary
        if n_bytes > 0:
            new_pos = self.pos_ + n_bytes
            memoryview(dest_buffer)[dest_pos:dest_pos+n_bytes] = \
                self.buf_[self.pos_:new_pos]
            self.pos_ = new_pos