kaitaistruct.py

import itertools, sys, struct
from io import open, BytesIO, SEEK_CUR, SEEK_END  # noqa

# Kaitai Struct runtime streaming API version, defined as per PEP-0396 standard.
# Used for two purposes:
#
# * .py files generated by ksc from .ksy check that they import proper KS runtime 
# * library by this version number;
# * distribution utils (setup.py) use this when packaging for PyPI
#
__version__ = '0.10'


class KaitaiStruct(object):
    def __init__(self, _io, _parent=None, _root=None, _init=False):
        self._io = _io
        self._parent = _parent
        self._root = _root if _root else self
        self._init = _init

    def __enter__(self):
        return self

    def __exit__(self, *args, **kwargs):
        self.close()

    def close(self):
        self._io.close()

    # _init used to config KaitaiStruct-derived calss manually and call _read()
    @classmethod
    def from_file(cls, filename, mode='rb', _init=True):
        f = open(filename, mode)
        try:
            return cls(KaitaiStream(f), _init=_init)
        except Exception:
            # close file descriptor, then reraise the exception otherwise
            # close the stream manually
            f.close()
            raise

    @classmethod
    def from_bytes(cls, buf):
        return cls(KaitaiStream(BytesIO(buf)), _init=True)

    @classmethod
    def from_io(cls, io):
        return cls(KaitaiStream(io), _init=True)

    def to_file(self, filename, mode='wb'):
        self.close()
        if self._io:
            self._io._io = open(filename, mode)
        else:
            self._io = KaitaiStream(open(filename, mode))
        try:
            self._write()
        except Exception:
            # close file descriptor, then reraise the exception otherwise
            # close the stream manually
            self.close()
            raise

    def to_bytes(self, pre=b''):
        self.close()
        if self._io:
            self._io._io = BytesIO(pre)
        else:
            self._io = KaitaiStream(BytesIO(pre))
        self._write()
        self._io.seek(0)
        return self._io.read_bytes_full()


class KaitaiStream(object):
    def __init__(self, io):
        self._io = io
        self.align_to_byte()

    def __enter__(self):
        return self

    def __exit__(self, *args, **kwargs):
        self.close()

    def close(self):
        self._io.close()

    # ========================================================================
    # Stream positioning
    # ========================================================================

    def is_eof(self):
        if self.bits_left > 0:
            return False

        io = self._io
        t = io.read(1)
        if t == b'':
            return True
        else:
            io.seek(-1, SEEK_CUR)
            return False

    def seek(self, n):
        self._io.seek(n)

    def pos(self):
        return self._io.tell()

    def size(self):
        # Python has no internal File object API function to get
        # current file / StringIO size, thus we use the following
        # trick.
        io = self._io
        # Remember our current position
        cur_pos = io.tell()
        # Seek to the end of the File object
        io.seek(0, SEEK_END)
        # Remember position, which is equal to the full length
        full_size = io.tell()
        # Seek back to the current position
        io.seek(cur_pos)
        return full_size

    # ========================================================================
    # Integer numbers
    # ========================================================================

    packer_s1 = struct.Struct('b')
    packer_s2be = struct.Struct('>h')
    packer_s4be = struct.Struct('>i')
    packer_s8be = struct.Struct('>q')
    packer_s2le = struct.Struct('<h')
    packer_s4le = struct.Struct('<i')
    packer_s8le = struct.Struct('<q')

    packer_u1 = struct.Struct('B')
    packer_u2be = struct.Struct('>H')
    packer_u4be = struct.Struct('>I')
    packer_u8be = struct.Struct('>Q')
    packer_u2le = struct.Struct('<H')
    packer_u4le = struct.Struct('<I')
    packer_u8le = struct.Struct('<Q')

    # ------------------------------------------------------------------------
    # Signed
    # ------------------------------------------------------------------------

    def read_s1(self):
        return KaitaiStream.packer_s1.unpack(self.read_bytes(1))[0]

    # ........................................................................
    # Big-endian
    # ........................................................................

    def read_s2be(self):
        return KaitaiStream.packer_s2be.unpack(self.read_bytes(2))[0]

    def read_s4be(self):
        return KaitaiStream.packer_s4be.unpack(self.read_bytes(4))[0]

    def read_s8be(self):
        return KaitaiStream.packer_s8be.unpack(self.read_bytes(8))[0]

    # ........................................................................
    # Little-endian
    # ........................................................................

    def read_s2le(self):
        return KaitaiStream.packer_s2le.unpack(self.read_bytes(2))[0]

    def read_s4le(self):
        return KaitaiStream.packer_s4le.unpack(self.read_bytes(4))[0]

    def read_s8le(self):
        return KaitaiStream.packer_s8le.unpack(self.read_bytes(8))[0]

    # ------------------------------------------------------------------------
    # Unsigned
    # ------------------------------------------------------------------------

    def read_u1(self):
        return KaitaiStream.packer_u1.unpack(self.read_bytes(1))[0]

    # ........................................................................
    # Big-endian
    # ........................................................................

    def read_u2be(self):
        return KaitaiStream.packer_u2be.unpack(self.read_bytes(2))[0]

    def read_u4be(self):
        return KaitaiStream.packer_u4be.unpack(self.read_bytes(4))[0]

    def read_u8be(self):
        return KaitaiStream.packer_u8be.unpack(self.read_bytes(8))[0]

    # ........................................................................
    # Little-endian
    # ........................................................................

    def read_u2le(self):
        return KaitaiStream.packer_u2le.unpack(self.read_bytes(2))[0]

    def read_u4le(self):
        return KaitaiStream.packer_u4le.unpack(self.read_bytes(4))[0]

    def read_u8le(self):
        return KaitaiStream.packer_u8le.unpack(self.read_bytes(8))[0]

    # ========================================================================
    # Floating point numbers
    # ========================================================================

    packer_f4be = struct.Struct('>f')
    packer_f8be = struct.Struct('>d')
    packer_f4le = struct.Struct('<f')
    packer_f8le = struct.Struct('<d')

    # ........................................................................
    # Big-endian
    # ........................................................................

    def read_f4be(self):
        return KaitaiStream.packer_f4be.unpack(self.read_bytes(4))[0]

    def read_f8be(self):
        return KaitaiStream.packer_f8be.unpack(self.read_bytes(8))[0]

    # ........................................................................
    # Little-endian
    # ........................................................................

    def read_f4le(self):
        return KaitaiStream.packer_f4le.unpack(self.read_bytes(4))[0]

    def read_f8le(self):
        return KaitaiStream.packer_f8le.unpack(self.read_bytes(8))[0]

    # ------------------------------------------------------------------------
    # Writing
    # ------------------------------------------------------------------------
    # ------------------------------------------------------------------------
    # Signed
    # ------------------------------------------------------------------------

    def write_s1(self, data):
        return self.write_bytes(KaitaiStream.packer_s1.pack(data))

    # ........................................................................
    # Big-endian
    # ........................................................................

    def write_s2be(self, data):
        return self.write_bytes(KaitaiStream.packer_s2be.pack(data))

    def write_s4be(self, data):
        return self.write_bytes(KaitaiStream.packer_s4be.pack(data))

    def write_s8be(self, data):
        return self.write_bytes(KaitaiStream.packer_s8be.pack(data))

    # ........................................................................
    # Little-endian
    # ........................................................................

    def write_s2le(self, data):
        return self.write_bytes(KaitaiStream.packer_s2le.pack(data))

    def write_s4le(self, data):
        return self.write_bytes(KaitaiStream.packer_s4le.pack(data))

    def write_s8le(self, data):
        return self.write_bytes(KaitaiStream.packer_s8le.pack(data))

    # ------------------------------------------------------------------------
    # Unsigned
    # ------------------------------------------------------------------------

    def write_u1(self, data):
        return self.write_bytes(KaitaiStream.packer_u1.pack(data))

    # ........................................................................
    # Big-endian
    # ........................................................................

    def write_u2be(self, data):
        return self.write_bytes(KaitaiStream.packer_u2be.pack(data))

    def write_u4be(self, data):
        return self.write_bytes(KaitaiStream.packer_u4be.pack(data))

    def write_u8be(self, data):
        return self.write_bytes(KaitaiStream.packer_u8be.pack(data))

    # ........................................................................
    # Little-endian
    # ........................................................................

    def write_u2le(self, data):
        return self.write_bytes(KaitaiStream.packer_u2le.pack(data))

    def write_u4le(self, data):
        return self.write_bytes(KaitaiStream.packer_u4le.pack(data))

    def write_u8le(self, data):
        return self.write_bytes(KaitaiStream.packer_u8le.pack(data))

    # ........................................................................
    # Big-endian
    # ........................................................................

    def write_f4be(self, data):
        return self.write_bytes(KaitaiStream.packer_f4be.pack(data))

    def write_f8be(self, data):
        return self.write_bytes(KaitaiStream.packer_f8be.pack(data))

    # ........................................................................
    # Little-endian
    # ........................................................................

    def write_f4le(self, data):
        return self.write_bytes(KaitaiStream.packer_f4le.pack(data))

    def write_f8le(self, data):
        return self.write_bytes(KaitaiStream.packer_f8le.pack(data))

    # ........................................................................
    # Ruby VLB integer
    # ........................................................................
    def read_ruby_long(self):
        length = self.read_s1()
        if length == 0:
            return 0
        if 5 < length < 128:
            return length - 5
        elif -129 < length < -5:
            return length + 5
        result = 0
        factor = 1
        for s in range(abs(length)):
            result += self.read_u1() * factor
            factor *= 256
        if length < 0:
            result = result - factor
        return result

    def write_ruby_long(self, data):
        if data == 0:
            self.write_u1(0)
        elif 0 < data < 123:
            self.write_s1(data + 5)
        elif -124 < data < 0:
            self.write_s1(data - 5)
        else:
            size = int(math.ceil(data.bit_length() / 8.0))
            if size > 5:
                raise ValueError("%d too long for serialization" % data)
            original_data = data
            factor = 256 ** size
            if data < 0 and data == -factor:
                size -= 1
                data += factor / 256
            elif data < 0:
                data += factor
            sign = int(math.copysign(size, original_data))
            self.write_s1(sign)
            for i in range(size):
                self.write_u1(data % 256)
                data //= 256

    # ========================================================================
    # Unaligned bit values
    # ========================================================================

    def align_to_byte(self):
        self.bits = 0
        self.bits_left = 0

    def read_bits_int_be(self, n):
        bits_needed = n - self.bits_left
        if bits_needed > 0:
            # 1 bit  => 1 byte
            # 8 bits => 1 byte
            # 9 bits => 2 bytes
            bytes_needed = ((bits_needed - 1) // 8) + 1
            buf = self.read_bytes(bytes_needed)
            for byte in buf:
                byte = KaitaiStream.int_from_byte(byte)
                self.bits <<= 8
                self.bits |= byte
                self.bits_left += 8

        # raw mask with required number of 1s, starting from lowest bit
        mask = (1 << n) - 1
        # shift self.bits to align the highest bits with the mask & derive reading result
        shift_bits = self.bits_left - n
        res = (self.bits >> shift_bits) & mask
        # clear top bits that we've just read => AND with 1s
        self.bits_left -= n
        mask = (1 << self.bits_left) - 1
        self.bits &= mask

        return res

    # Unused since Kaitai Struct Compiler v0.9+ - compatibility with
    # older versions.
    def read_bits_int(self, n):
        return self.read_bits_int_be(n)

    def read_bits_int_le(self, n):
        bits_needed = n - self.bits_left
        if bits_needed > 0:
            # 1 bit  => 1 byte
            # 8 bits => 1 byte
            # 9 bits => 2 bytes
            bytes_needed = ((bits_needed - 1) // 8) + 1
            buf = self.read_bytes(bytes_needed)
            for byte in buf:
                byte = KaitaiStream.int_from_byte(byte)
                self.bits |= (byte << self.bits_left)
                self.bits_left += 8

        # raw mask with required number of 1s, starting from lowest bit
        mask = (1 << n) - 1
        # derive reading result
        res = self.bits & mask
        # remove bottom bits that we've just read by shifting
        self.bits >>= n
        self.bits_left -= n

        return res

    # ========================================================================
    # Byte arrays
    # ========================================================================

    def alignment(self, a):
        return (a - self.pos()) % a

    def read_bytes(self, n):
        if n < 0:
            raise ValueError(
                "requested invalid %d amount of bytes" %
                (n,)
            )
        r = self._io.read(n)
        if len(r) < n:
            raise EOFError(
                "requested %d bytes, but got only %d bytes" %
                (n, len(r))
            )
        return r

    def seek_alignment(self, align):
        if align > 1:
            self._io.seek(self.alignment(align), 1)

    def write_bytes(self, data):
        if data is None:
            return
        nb = len(data)
        if nb > 0 and self._io.write(data) != nb:
            raise Exception(
                "not all of %d bytes written" %
                (nb,)
            )
        return nb

    def write_alignment(self, align, align_byte=0):
        return self._io.write(KaitaiStream.packer_u1.pack(align_byte) * self.alignment(align))

    def write_terminator(self, term_bytes):
        return self._io.write(bytes(term_bytes))

    def write_padding(self, actual_size, size, pad_bytes):
        pad = size - actual_size
        if pad > 0:
            if pad_byte is not None:
                return self._io.write(bytes(pad_bytes) * (pad // len(pad_bytes) + pad % len(pad_bytes)))
            else:
                raise Exception("no padding filler provided")
        return 0

    def read_bytes_full(self):
        return self._io.read()

    def read_bytes_term(self, term_bytes, include_term, consume_term, eos_error, elem_size=1):
        r = b''
        term_b = bytes(term_bytes)
        while True:
            c = self._io.read(elem_size)
            if len(c) != elem_size:
                if eos_error:
                    raise EOFError(
                        "end of stream reached, but no terminator %d found" %
                        (term,)
                    )
                else:
                    return r
            elif c == term_b:
                if include_term:
                    r += c
                if not consume_term:
                    self._io.seek(-elem_size, SEEK_CUR)
                return r
            else:
                r += c

    def ensure_fixed_contents(self, expected):
        actual = self._io.read(len(expected))
        if actual != expected:
            raise Exception(
                "unexpected fixed contents: got %r, was waiting for %r" %
                (actual, expected)
            )
        return actual

    @staticmethod
    def bytes_strip_right(data, pad_bytes, elem_size=1):
        new_len = len(data)
        pad_bytes = tuple(pad_bytes)
        while new_len > 0 and tuple(data[new_len-elem_size:new_len]) == pad_bytes:
            new_len -= elem_size

        return data[:new_len]

    @staticmethod
    def bytes_terminate(data, term, include_term, elem_size=1):
        new_len = 0
        max_len = len(data)
        term =  tuple(term)

        while new_len < max_len and tuple(data[new_len:new_len+elem_size]) != term:
            new_len += elem_size

        if include_term and new_len < max_len:
            new_len += elem_size

        return data[:new_len]

    # ========================================================================
    # Byte array processing
    # ========================================================================

    @staticmethod
    def process_xor_one(data, key):
         return bytes(v ^ key for v in data)

    @staticmethod
    def process_xor_many(data, key):
        return bytes(a ^ b for a, b in zip(data, itertools.cycle(key)))

    @staticmethod
    def process_rotate_left(data, amount, group_size):
        if group_size != 1:
            raise Exception(
                "unable to rotate group of %d bytes yet" %
                (group_size,)
            )

        mask = group_size * 8 - 1
        anti_amount = -amount & mask

        r = bytearray(data)
        for i in range(len(r)):
            r[i] = (r[i] << amount) & 0xff | (r[i] >> anti_amount)
        return bytes(r)

    # ========================================================================
    # Misc
    # ========================================================================

    @staticmethod
    def int_from_byte(v):
        return v

    @staticmethod
    def byte_array_index(data, i):
        return KaitaiStream.int_from_byte(data[i])

    @staticmethod
    def byte_array_min(b):
        return KaitaiStream.int_from_byte(min(b))

    @staticmethod
    def byte_array_max(b):
        return KaitaiStream.int_from_byte(max(b))

    @staticmethod
    def resolve_enum(enum_obj, value):
        """Resolves value using enum: if the value is not found in the map,
        we'll just use literal value per se. Works around problem with Python
        enums throwing an exception when encountering unknown value.
        """
        try:
            return enum_obj(value)
        except ValueError:
            return value


class KaitaiStructError(BaseException):
    """Common ancestor for all error originating from Kaitai Struct usage.
    Stores KSY source path, pointing to an element supposedly guilty of
    an error.
    """
    def __init__(self, msg, src_path):
        super(KaitaiStructError, self).__init__("%s: %s" % (src_path, msg))
        self.src_path = src_path


class UndecidedEndiannessError(KaitaiStructError):
    """Error that occurs when default endianness should be decided with
    switch, but nothing matches (although using endianness expression
    implies that there should be some positive result).
    """
    def __init__(self, src_path):
        super(KaitaiStructError, self).__init__("unable to decide on endianness for a type", src_path)


class ValidationFailedError(KaitaiStructError):
    """Common ancestor for all validation failures. Stores pointer to
    KaitaiStream IO object which was involved in an error.
    """
    def __init__(self, msg, io, src_path):
        super(ValidationFailedError, self).__init__("at pos %d: validation failed: %s" % (io.pos(), msg), src_path)
        self.io = io


class ValidationNotEqualError(ValidationFailedError):
    """Signals validation failure: we required "actual" value to be equal to
    "expected", but it turned out that it's not.
    """
    def __init__(self, expected, actual, io, src_path):
        super(ValidationNotEqualError, self).__init__("not equal, expected %s, but got %s" % (repr(expected), repr(actual)), io, src_path)
        self.expected = expected
        self.actual = actual


class ValidationLessThanError(ValidationFailedError):
    """Signals validation failure: we required "actual" value to be
    greater than or equal to "min", but it turned out that it's not.
    """
    def __init__(self, _min, actual, io, src_path):
        super(ValidationLessThanError, self).__init__("not in range, min %s, but got %s" % (repr(min), repr(actual)), io, src_path)
        self.min = _min
        self.actual = actual


class ValidationGreaterThanError(ValidationFailedError):
    """Signals validation failure: we required "actual" value to be
    less than or equal to "max", but it turned out that it's not.
    """
    def __init__(self, _max, actual, io, src_path):
        super(ValidationGreaterThanError, self).__init__("not in range, max %s, but got %s" % (repr(max), repr(actual)), io, src_path)
        self.max = _max
        self.actual = actual


class ValidationNotAnyOfError(ValidationFailedError):
    """Signals validation failure: we required "actual" value to be
    from the list, but it turned out that it's not.
    """
    def __init__(self, actual, io, src_path):
        super(ValidationNotAnyOfError, self).__init__("not any of the list, got %s" % (repr(actual)), io, src_path)
        self.actual = actual


class ValidationExprError(ValidationFailedError):
    """Signals validation failure: we required "actual" value to match
    the expression, but it turned out that it doesn't.
    """
    def __init__(self, actual, io, src_path):
        super(ValidationExprError, self).__init__("not matching the expression, got %s" % (repr(actual)), io, src_path)
        self.actual = actual