cstruct.py

# Copyright (c) 2018 Fox-IT Security Research Team <srt@fox-it.com>
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

# TODO:
# - Rework definition parsing, maybe pycparser?
# - Change expression implementation
# - Lazy reading?
from __future__ import print_function
import re
import sys
import ast
import pprint
import string
import struct
import ctypes as _ctypes
from io import BytesIO
from collections import OrderedDict

try:
    from builtins import bytes as newbytes
except ImportError:
    newbytes = bytes

PY3 = sys.version_info > (3,)
if PY3:
    long = int
    xrange = range

DEBUG = False

COLOR_RED = '\033[1;31m'
COLOR_GREEN = '\033[1;32m'
COLOR_YELLOW = '\033[1;33m'
COLOR_BLUE = '\033[1;34m'
COLOR_PURPLE = '\033[1;35m'
COLOR_CYAN = '\033[1;36m'
COLOR_WHITE = '\033[1;37m'
COLOR_NORMAL = '\033[1;0m'

COLOR_BG_RED = '\033[1;41m\033[1;37m'
COLOR_BG_GREEN = '\033[1;42m\033[1;37m'
COLOR_BG_YELLOW = '\033[1;43m\033[1;37m'
COLOR_BG_BLUE = '\033[1;44m\033[1;37m'
COLOR_BG_PURPLE = '\033[1;45m\033[1;37m'
COLOR_BG_CYAN = '\033[1;46m\033[1;37m'
COLOR_BG_WHITE = '\033[1;47m\033[1;30m'

PRINTABLE = string.digits + string.ascii_letters + string.punctuation + " "

COMPILE_TEMPL = """
class {name}(Structure):
    def __init__(self, cstruct, structure, source=None):
        self.structure = structure
        self.source = source
        super({name}, self).__init__(cstruct, structure.name, structure.fields)

    def _read(self, stream):
        r = OrderedDict()
        sizes = {{}}
        bitreader = BitBuffer(stream, self.cstruct.endian)

{read_code}

        return Instance(self, r, sizes)

    def add_fields(self, name, type_, offset=None):
        raise NotImplementedError("Can't add fields to a compiled structure")

    def __repr__(self):
        return '<Structure {name} +compiled>'
"""


class Error(Exception):
    pass


class ParserError(Error):
    pass


class CompilerError(Error):
    pass


class ResolveError(Error):
    pass


class NullPointerDereference(Error):
    pass


def log(line, *args, **kwargs):
    if not DEBUG:
        return

    print(line.format(*args, **kwargs), file=sys.stderr)


class cstruct(object):
    """Main class of cstruct. All types are registered in here.

    Args:
        endian: The endianness to use when parsing.
        pointer: The pointer type to use for Pointers.
    """

    DEF_CSTYLE = 1

    def __init__(self, endian='<', pointer='uint64'):
        self.endian = endian

        self.consts = {}
        self.lookups = {}
        self.typedefs = {
            'byte': 'int8',
            'ubyte': 'uint8',
            'uchar': 'uint8',
            'short': 'int16',
            'ushort': 'uint16',
            'long': 'int32',
            'ulong': 'uint32',
            'ulong64': 'uint64',

            'u1': 'uint8',
            'u2': 'uint16',
            'u4': 'uint32',
            'u8': 'uint64',

            'word': 'uint16',
            'dword': 'uint32',

            'longlong': 'int64',
            'ulonglong': 'uint64',

            'int': 'int32',
            'unsigned int': 'uint32',

            'int8': PackedType(self, 'int8', 1, 'b'),
            'uint8': PackedType(self, 'uint8', 1, 'B'),
            'int16': PackedType(self, 'int16', 2, 'h'),
            'uint16': PackedType(self, 'uint16', 2, 'H'),
            'int32': PackedType(self, 'int32', 4, 'i'),
            'uint32': PackedType(self, 'uint32', 4, 'I'),
            'int64': PackedType(self, 'int64', 8, 'q'),
            'uint64': PackedType(self, 'uint64', 8, 'Q'),
            'float': PackedType(self, 'float', 4, 'f'),
            'double': PackedType(self, 'double', 8, 'd'),
            'char': CharType(self),
            'wchar': WcharType(self),

            'int24': BytesInteger(self, 'int24', 3, True),
            'uint24': BytesInteger(self, 'uint24', 3, False),
            'int48': BytesInteger(self, 'int48', 6, True),
            'uint48': BytesInteger(self, 'uint48', 6, False),

            'void': VoidType(),
        }

        self.pointer = self.resolve(pointer)

    def addtype(self, name, t, replace=False):
        """Add a type or type reference.

        Args:
            name: Name of the type to be added.
            t: The type to be added. Can be a str reference to another type
                or a compatible type class.

        Raises:
            ValueError: If the type already exists.
        """
        name = name.lower()
        if not replace and name.lower() in self.typedefs:
            raise ValueError("Duplicate type: %s" % name)

        self.typedefs[name] = t

    def load(self, s, deftype=None, **kwargs):
        """Parse structures from the given definitions using the given definition type.

        Definitions can be parsed using different parsers. Currently, there's
        only one supported parser - DEF_CSTYLE. Parsers can add types and
        modify this cstruct instance. Arguments can be passed to parsers
        using kwargs.

        Args:
            s: The definition to parse.
            deftype: The definition type to parse the definitions with.
            **kwargs: Keyword arguments for parsers.
        """
        deftype = deftype or cstruct.DEF_CSTYLE

        if deftype == cstruct.DEF_CSTYLE:
            parser = CStyleParser(self, **kwargs)

        parser.parse(s)

    def loadfile(self, s, deftype=None, **kwargs):
        """Load structure definitions from a file.

        The given path will be read and parsed using the .load() function.

        Args:
            s: The path to load definitions from.
            deftype: The definition type to parse the definitions with.
            **kwargs: Keyword arguments for parsers.
        """
        with open(s, 'r') as fh:
            self.load(fh.read(), deftype, **kwargs)

    def read(self, name, s):
        """Parse data using a given type.

        Args:
            name: Type name to read.
            s: File-like object or byte string to parse.

        Returns:
            The parsed data.
        """
        return self.resolve(name).read(s)

    def resolve(self, name):
        """Resolve a type name to get the actual type object.

        Types can be referenced using different names. When we want
        the actual type object, we need to resolve these references.

        Args:
            name: Type name to resolve.

        Returns:
            The resolved type object.

        Raises:
            ResolveError: If the type can't be resolved.
        """
        t = name
        if not isinstance(t, str):
            return t

        for i in xrange(10):
            if t.lower() not in self.typedefs:
                raise ResolveError("Unknown type %s" % name)

            t = self.typedefs[t.lower()]

            if not isinstance(t, str):
                return t

        raise ResolveError("Recursion limit exceeded while resolving type %s" % name)

    def __getattr__(self, attr):
        if attr.lower() in self.typedefs:
            return self.typedefs[attr.lower()]

        if attr in self.consts:
            return self.consts[attr]

        raise AttributeError("Invalid Attribute: %s" % attr)


class Parser(object):
    """Base class for definition parsers.

    Args:
        cstruct: An instance of cstruct.
    """

    def __init__(self, cstruct):
        self.cstruct = cstruct

    def parse(self, data):
        """This function should parse definitions to cstruct types.

        Args:
            data: Data to parse definitions from, usually a string.
        """
        raise NotImplementedError()


class CStyleParser(Parser):
    """Definition parser for C-like structure syntax.

    Args:
        cstruct: An instance of cstruct
        compiled: Whether structs should be compiled or not.
    """

    def __init__(self, cstruct, compiled=True):
        self.compiled = compiled
        super(CStyleParser, self).__init__(cstruct)

    # TODO: Implement proper parsing
    def parse(self, data):
        self._constants(data)
        self._enums(data)
        self._structs(data)
        self._lookups(data, self.cstruct.consts)

    def _constants(self, data):
        r = re.finditer(r'#define\s+(?P<name>[^\s]+)\s+(?P<value>[^\r\n]+)\s*\n', data)
        for t in r:
            d = t.groupdict()
            v = d['value'].rsplit('//')[0]

            try:
                v = ast.literal_eval(v)
            except (ValueError, SyntaxError):
                pass

            self.cstruct.consts[d['name']] = v

    def _enums(self, data):
        r = re.finditer(
            r'enum\s+(?P<name>[^\s:{]+)\s*(:\s*(?P<type>[^\s]+)\s*)?\{(?P<values>[^}]+)\}\s*;',
            data,
        )
        for t in r:
            d = t.groupdict()

            nextval = 0
            values = {}
            for line in d['values'].split('\n'):
                line, sep, comment = line.partition("//")
                for v in line.split(","):
                    key, sep, val = v.partition("=")
                    key = key.strip()
                    val = val.strip()
                    if not key:
                        continue
                    if not val:
                        val = nextval
                    else:
                        val = Expression(self.cstruct, val).evaluate({})

                    nextval = val + 1

                    values[key] = val

            if not d['type']:
                d['type'] = 'uint32'

            enum = Enum(
                self.cstruct, d['name'], self.cstruct.resolve(d['type']), values
            )
            self.cstruct.addtype(enum.name, enum)

    def _structs(self, data):
        compiler = Compiler(self.cstruct)
        r = re.finditer(
            r'(#(?P<flags>(?:compile))\s+)?((?P<typedef>typedef)\s+)?(?P<type>[^\s]+)\s+(__attribute__\(\([^)]+\)\)\s*)?(?P<name>[^\s]+)?(?P<fields>\s*\{[^}]+\}(?P<defs>\s+[^;\n]+)?)?\s*;',
            data,
        )
        for t in r:
            d = t.groupdict()

            if d['name']:
                name = d['name']
            elif d['defs']:
                name = d['defs'].strip().split(',')[0].strip()
            else:
                raise ParserError("No name for struct")

            if d['type'] == 'struct':
                data = self._parse_fields(d['fields'][1:-1].strip())
                st = Structure(self.cstruct, name, data)
                if d['flags'] == 'compile' or self.compiled:
                    st = compiler.compile(st)
            elif d['typedef'] == 'typedef':
                st = d['type']
            else:
                continue

            if d['name']:
                self.cstruct.addtype(d['name'], st)

            if d['defs']:
                for td in d['defs'].strip().split(','):
                    td = td.strip()
                    self.cstruct.addtype(td, st)

    def _parse_fields(self, s):
        fields = re.finditer(
            r'(?P<type>[^\s]+)\s+(?P<name>[^\s\[:]+)(\s*:\s*(?P<bits>\d+))?(\[(?P<count>[^;\n]*)\])?;',
            s,
        )
        r = []
        for f in fields:
            d = f.groupdict()
            if d['type'].startswith('//'):
                continue

            type_ = self.cstruct.resolve(d['type'])

            d['name'] = d['name'].replace('(', '').replace(')', '')

            # Maybe reimplement lazy type references later
            # _type = TypeReference(self, d['type'])
            if d['count'] is not None:
                if d['count'] == '':
                    count = None
                else:
                    count = Expression(self.cstruct, d['count'])
                    try:
                        count = count.evaluate()
                    except Exception:
                        pass

                type_ = Array(self.cstruct, type_, count)

            if d['name'].startswith('*'):
                d['name'] = d['name'][1:]
                type_ = Pointer(self.cstruct, type_)

            field = Field(d['name'], type_, int(d['bits']) if d['bits'] else None)
            r.append(field)

        return r

    def _lookups(self, data, consts):
        r = re.finditer(r'\$(?P<name>[^\s]+) = ({[^}]+})\w*\n', data)

        for t in r:
            d = ast.literal_eval(t.group(2))
            self.cstruct.lookups[t.group(1)] = dict(
                [(self.cstruct.consts[k], v) for k, v in d.items()]
            )


class Instance(object):
    """Holds parsed structure data."""

    def __init__(self, type_, values, sizes=None):
        object.__setattr__(self, '_type', type_)
        object.__setattr__(self, '_values', values)
        object.__setattr__(self, '_sizes', sizes)

    def write(self, fh):
        """Write this structure to a writable file-like object.

        Args:
            fh: File-like objects that supports writing.

        Returns:
            The amount of bytes written.
        """
        return self.__dict__['_type'].write(fh, self)

    def dumps(self):
        """Dump this structure to a byte string.

        Returns:
            The raw bytes of this structure.
        """
        s = BytesIO()
        self.write(s)
        return s.getvalue()

    def __getattr__(self, attr):
        if attr not in self.__dict__['_type'].lookup:
            raise AttributeError("Invalid attribute: %r" % attr)

        return self.__dict__['_values'][attr]

    def __setattr__(self, attr, value):
        if attr not in self.__dict__['_type'].lookup:
            raise AttributeError("Invalid attribute: %r" % attr)

        self.__dict__['_values'][attr] = value

    def __getitem__(self, item):
        return self.__dict__['_values'][item]

    def __contains__(self, attr):
        return attr in self.__dict__['_values']

    def __repr__(self):
        return '<%s %s>' % (
            self.__dict__['_type'].name,
            ', '.join(
                [
                    '%s=%s' % (k, hex(v) if isinstance(v, (int, long)) else repr(v))
                    for k, v in self.__dict__['_values'].items()
                ]
            ),
        )

    def __len__(self):
        return len(self.dumps())

    def _size(self, field):
        return self.__dict__['_sizes'][field]


class PointerInstance(object):
    """Like the Instance class, but for structures referenced by a pointer."""

    def __init__(self, t, stream, addr, ctx):
        self._stream = stream
        self._type = t
        self._addr = addr
        self._ctx = ctx
        self._value = None

    def _get(self):
        log("Dereferencing pointer -> 0x{:016x} [{!r}]", self._addr, self._stream)
        if self._addr == 0:
            raise NullPointerDereference()

        if self._value is None:
            pos = self._stream.tell()
            self._stream.seek(self._addr)
            if isinstance(self._type, Array):
                r = self._type._read(self._stream, self._ctx)
            else:
                r = self._type._read(self._stream)
            self._stream.seek(pos)
            self._value = r

        return self._value

    def __getattr__(self, attr):
        return getattr(self._get(), attr)

    def __str__(self):
        return str(self._get())

    def __nonzero__(self):
        return self._addr != 0

    def __repr__(self):
        return "<Pointer {!r} @ 0x{:x}>".format(self._type, self._addr)


class Expression(object):
    """Expression parser for simple calculations in definitions."""

    operators = [
        ('+', lambda a, b: a + b),
        ('-', lambda a, b: a - b),
        ('*', lambda a, b: a * b),
        ('/', lambda a, b: a / b),
        ('&', lambda a, b: a & b),
        ('|', lambda a, b: a | b),
        ('>>', lambda a, b: a >> b),
        ('<<', lambda a, b: a << b),
    ]

    def __init__(self, cstruct, expr):
        self.cstruct = cstruct
        self.expr = expr

    def evaluate(self, context=None):
        context = context if context else {}
        level = 0
        levels = []
        buf = ''

        for i in xrange(len(self.expr)):
            if self.expr[i] == '(':
                level += 1
                levels.append(buf)
                buf = ''
                continue

            if self.expr[i] == ')':
                level -= 1
                val = self.evaluate_part(buf, context)
                buf = levels.pop()
                buf += str(val)
                continue

            buf += self.expr[i]

        return self.evaluate_part(buf, context)

    def evaluate_part(self, e, v):
        e = e.strip()

        for o in self.operators:
            if o[0] in e:
                a, b = e.rsplit(o[0], 1)
                return o[1](self.evaluate_part(a, v), self.evaluate_part(b, v))

        if e in v:
            return v[e]

        if e.startswith('0x'):
            return int(e, 16)

        if e in self.cstruct.consts:
            return self.cstruct.consts[e]

        return int(e)

    def __repr__(self):
        return self.expr


class BaseType(object):
    """Base class for cstruct type classes."""

    def __init__(self, cstruct):
        self.cstruct = cstruct

    def reads(self, data):
        """Parse the given data according to the type that implements this class.

        Args:
            data: Byte string to parse.

        Returns:
            The parsed value of this type.
        """
        data = BytesIO(data)
        return self._read(data)

    def dumps(self, data):
        """Dump the given data according to the type that implements this class.

        Args:
            data: Data to dump.

        Returns:
            The resulting bytes.
        """
        out = BytesIO()
        self._write(out, data)
        return out.getvalue()

    def read(self, obj, *args, **kwargs):
        """Parse the given data according to the type that implements this class.

        Args:
            obj: Data to parse. Can be a (byte) string or a file-like object.

        Returns:
            The parsed value of this type.
        """
        if isinstance(obj, (str, bytes, newbytes)):
            return self.reads(obj)

        return self._read(obj)

    def write(self, stream, data):
        """Write the given data to a writable file-like object according to the
        type that implements this class.

        Args:
            stream: Writable file-like object to write to.
            data: Data to write.

        Returns:
            The amount of bytes written.
        """
        return self._write(stream, data)

    def _read(self, stream):
        raise NotImplementedError()

    def _read_array(self, stream, count):
        return [self._read(stream) for i in xrange(count)]

    def _read_0(self, stream):
        raise NotImplementedError()

    def _write(self, stream, data):
        raise NotImplementedError()

    def _write_array(self, stream, data):
        num = 0
        for i in data:
            num += self._write(stream, i)
        return num

    def _write_0(self, stream, data):
        raise NotImplementedError()

    def default(self):
        """Return a default value of this type."""
        raise NotImplementedError()

    def default_array(self):
        """Return a default array of this type."""
        raise NotImplementedError()

    def __getitem__(self, count):
        return Array(self.cstruct, self, count)

    def __call__(self, *args, **kwargs):
        if len(args) > 0:
            return self.read(*args, **kwargs)

        r = self.default()
        if kwargs:
            for k, v in kwargs.items():
                setattr(r, k, v)

        return r


class RawType(BaseType):
    """Base class for raw types that have a name and size."""

    def __init__(self, cstruct, name=None, size=0):
        self.name = name
        self.size = size
        super(RawType, self).__init__(cstruct)

    def __len__(self):
        return self.size

    def __repr__(self):
        if self.name:
            return self.name

        return BaseType.__repr__(self)


class Structure(BaseType):
    """Type class for structures."""

    def __init__(self, cstruct, name, fields=None):
        self.name = name
        self.size = None
        self.lookup = OrderedDict()
        self.fields = fields if fields else []

        for f in self.fields:
            self.lookup[f.name] = f

        self._calc_offsets()
        super(Structure, self).__init__(cstruct)

    def _calc_offsets(self):
        offset = 0
        bitstype = None
        bitsremaining = 0

        for field in self.fields:
            if field.bits:
                if bitsremaining == 0 or field.type != bitstype:
                    bitstype = field.type
                    bitsremaining = bitstype.size * 8
                    if offset is not None:
                        field.offset = offset
                        offset += bitstype.size
                else:
                    field.offset = None

                bitsremaining -= field.bits
                continue

            field.offset = offset
            if offset is not None:
                try:
                    offset += len(field.type)
                except TypeError:
                    offset = None

    def _calc_size(self):
        size = 0
        bitstype = None
        bitsremaining = 0

        for field in self.fields:
            if field.bits:
                if bitsremaining == 0 or field.type != bitstype:
                    bitstype = field.type
                    bitsremaining = bitstype.size * 8
                    size += bitstype.size

                bitsremaining -= field.bits
                continue

            fieldlen = len(field.type)
            size += fieldlen

            if field.offset is not None:
                size = max(size, field.offset + fieldlen)

        return size

    def _read(self, stream, *args, **kwargs):
        log("[Structure::read] {} {}", self.name, self.size)
        bitbuffer = BitBuffer(stream, self.cstruct.endian)

        struct_start = stream.tell()

        r = OrderedDict()
        sizes = {}
        for field in self.fields:
            start = stream.tell()
            ft = self.cstruct.resolve(field.type)

            if field.offset:
                if start != struct_start + field.offset:
                    log(
                        "+ seeking to 0x{:x}+0x{:x} for {}".format(
                            struct_start, field.offset, field.name
                        )
                    )
                    stream.seek(struct_start + field.offset)
                    start = struct_start + field.offset

            if field.bits:
                r[field.name] = bitbuffer.read(ft, field.bits)
                continue
            else:
                bitbuffer.reset()

            if isinstance(ft, (Array, Pointer)):
                v = ft._read(stream, r)
            else:
                v = ft._read(stream)

            sizes[field.name] = stream.tell() - start
            r[field.name] = v

        return Instance(self, r, sizes)

    def _write(self, stream, data):
        bitbuffer = BitBuffer(stream, self.cstruct.endian)
        num = 0

        for field in self.fields:
            if field.bits:
                bitbuffer.write(field.type, getattr(data, field.name), field.bits)
                continue

            if bitbuffer._type:
                bitbuffer.flush()

            num += field.type._write(stream, getattr(data, field.name))

        # Flush bitbuffer
        if bitbuffer._type:
            bitbuffer.flush()

        return num

    def add_field(self, name, type_, offset=None):
        """Add a field to this structure.

        Args:
            name: The field name.
            type_: The field type.
            offset: The field offset.
        """
        field = Field(name, type_, offset=offset)
        self.fields.append(field)
        self.lookup[name] = field
        self.size = None
        setattr(self, name, field)

    def default(self):
        """Create and return an empty Instance from this structure.

        Returns:
            An empty Instance from this structure.
        """
        r = OrderedDict()
        for field in self.fields:
            r[field.name] = field.type.default()

        return Instance(self, r)

    def __len__(self):
        if self.size is None:
            self.size = self._calc_size()

        return self.size

    def __repr__(self):
        return '<Structure {}>'.format(self.name)

    def show(self, indent=0):
        """Pretty print this structure."""
        if indent == 0:
            print("struct {}".format(self.name))

        for field in self.fields:
            if field.offset is None:
                offset = '0x??'
            else:
                offset = '0x{:02x}'.format(field.offset)

            print("{}+{} {} {}".format(' ' * indent, offset, field.name, field.type))

            if isinstance(field.type, Structure):
                field.type.show(indent + 1)


class BitBuffer(object):
    """Implements a bit buffer that can read and write bit fields."""

    def __init__(self, stream, endian):
        self.stream = stream
        self.endian = endian

        self._type = None
        self._buffer = 0
        self._remaining = 0

    def read(self, field_type, bits):
        if self._remaining < 1 or self._type != field_type:
            self._type = field_type
            self._remaining = field_type.size * 8
            self._buffer = field_type._read(self.stream)

        if self.endian != '>':
            v = self._buffer & ((1 << bits) - 1)
            self._buffer >>= bits
            self._remaining -= bits
        else:
            v = self._buffer & (
                ((1 << (self._remaining - bits)) - 1) ^ ((1 << self._remaining) - 1)
            )
            v >>= self._remaining - bits
            self._remaining -= bits

        return v

    def write(self, field_type, data, bits):
        if self._remaining == 0:
            self._remaining = field_type.size * 8
            self._type = field_type

        if self.endian != '>':
            self._buffer |= data << (self._type.size * 8 - self._remaining)
        else:
            self._buffer |= data << (self._remaining - bits)

        self._remaining -= bits

    def flush(self):
        self._type._write(self.stream, self._buffer)
        self._type = None
        self._remaining = 0
        self._buffer = 0

    def reset(self):
        self._type = None
        self._buffer = 0
        self._remaining = 0


class Field(object):
    """Holds a structure field."""

    def __init__(self, name, type_, bits=None, offset=None):
        self.name = name
        self.type = type_
        self.bits = bits
        self.offset = offset

    def __repr__(self):
        return '<Field {} {}>'.format(self.name, self.type)


class Array(BaseType):
    """Implements a fixed or dynamically sized array type.

    Example:
        When using the default C-style parser, the following syntax is supported:

            x[3] -> 3 -> static length.
            x[] -> None -> null-terminated.
            x[expr] -> expr -> dynamic length.
    """

    def __init__(self, cstruct, type_, count):
        self.type = type_
        self.count = count
        self.dynamic = isinstance(self.count, Expression) or self.count is None

        super(Array, self).__init__(cstruct)

    def _read(self, stream, context=None):
        if self.count is None:
            return self.type._read_0(stream)

        if self.dynamic:
            count = self.count.evaluate(context)
        else:
            count = self.count

        return self.type._read_array(stream, max(0, count))

    def _write(self, f, data):
        if self.count is None:
            return self.type._write_0(f, data)

        return self.type._write_array(f, data)

    def default(self):
        if self.dynamic or self.count is None:
            return []

        return [self.type.default() for i in xrange(self.count)]

    def __repr__(self):
        if self.count is None:
            return '{0!r}[]'.format(self.type)

        return '{0!r}[{1}]'.format(self.type, self.count)

    def __len__(self):
        if self.dynamic:
            raise TypeError("Dynamic size")

        return len(self.type) * self.count


class PackedType(RawType):
    """Implements a packed type that uses Python struct packing characters."""

    def __init__(self, cstruct, name, size, packchar):
        self.packchar = packchar
        super(PackedType, self).__init__(cstruct, name, size)

    def _read(self, stream):
        return self._read_array(stream, 1)[0]

    def _read_array(self, stream, count):
        length = self.size * count
        data = stream.read(length)
        fmt = self.cstruct.endian + str(count) + self.packchar
        if len(data) != length:
            raise EOFError("Read %d bytes, but expected %d" % (len(data), length))

        return list(struct.unpack(fmt, data))

    def _read_0(self, stream):
        r = []
        while True:
            d = stream.read(self.size)
            v = struct.unpack(self.cstruct.endian + self.packchar, d)[0]

            if v == 0:
                break

            r.append(v)

        return r

    def _write(self, stream, data):
        return self._write_array(stream, [data])

    def _write_array(self, stream, data):
        fmt = self.cstruct.endian + str(len(data)) + self.packchar
        return stream.write(struct.pack(fmt, *data))

    def _write_0(self, stream, data):
        return self._write_array(stream, data + [0])

    def default(self):
        return 0

    def default_array(self, count):
        return [0] * count


class CharType(RawType):
    """Implements a character type that can properly handle strings."""

    def __init__(self, cstruct):
        super(CharType, self).__init__(cstruct, 'char', 1)

    def _read(self, stream):
        return stream.read(1)

    def _read_array(self, stream, count):
        if count == 0:
            return b''

        return stream.read(count)

    def _read_0(self, stream):
        r = []
        while True:
            c = stream.read(1)
            if c == b'':
                raise EOFError()

            if c == b'\x00':
                break

            r.append(c)

        return b''.join(r)

    def _write(self, stream, data):
        if isinstance(data, int):
            data = chr(data)

        if PY3 and isinstance(data, str):
            data = data.encode('latin-1')

        return stream.write(data)

    def _write_array(self, stream, data):
        return self._write(stream, data)

    def _write_0(self, stream, data):
        return self._write(stream, data + b'\x00')

    def default(self):
        return b'\x00'

    def default_array(self, count):
        return b'\x00' * count


class WcharType(RawType):
    """Implements a wide-character type."""

    def __init__(self, cstruct):
        super(WcharType, self).__init__(cstruct, 'wchar', 2)

    @property
    def encoding(self):
        if self.cstruct.endian == '<':
            return 'utf-16-le'
        elif self.cstruct.endian == '>':
            return 'utf-16-be'

    def _read(self, stream):
        return stream.read(2).decode(self.encoding)

    def _read_array(self, stream, count):
        if count == 0:
            return u''

        data = stream.read(2 * count)
        return data.decode(self.encoding)

    def _read_0(self, stream):
        r = b''
        while True:
            c = stream.read(2)

            if len(c) != 2:
                raise EOFError()

            if c == b'\x00\x00':
                break

            r += c

        return r.decode(self.encoding)

    def _write(self, stream, data):
        return stream.write(data.encode(self.encoding))

    def _write_array(self, stream, data):
        return self._write(stream, data)

    def _write_0(self, stream, data):
        return self._write(stream, data + u'\x00')

    def default(self):
        return u'\x00'

    def default_array(self, count):
        return u'\x00' * count


class BytesInteger(RawType):
    """Implements an integer type that can span an arbitrary amount of bytes."""

    def __init__(self, cstruct, name, size, signed):
        self.signed = signed
        super(BytesInteger, self).__init__(cstruct, name, size)

    @staticmethod
    def parse(buf, size, count, signed, endian):
        nums = []

        for c in xrange(count):
            num = 0
            data = buf[c * size:(c + 1) * size]
            if endian == '<':
                data = b''.join(data[i:i + 1] for i in reversed(xrange(len(data))))

            ints = list(data) if PY3 else map(ord, data)
            for i in ints:
                num = (num << 8) | i

            if signed and num & 1 << (size * 8 - 1):
                bias = 1 << (size * 8 - 1)
                num -= bias * 2

            nums.append(num)

        return nums

    @staticmethod
    def pack(data, size, endian):
        buf = []
        for i in data:
            num = int(i)
            if num < 0:
                num += 1 << (size * 8)

            d = [b'\x00'] * size
            i = size - 1

            while i >= 0:
                b = num & 255
                d[i] = bytes((b,)) if PY3 else chr(b)
                num >>= 8
                i -= 1

            if endian == '<':
                d = b''.join(d[i:i + 1][0] for i in reversed(xrange(len(d))))
            else:
                d = b''.join(d)

            buf.append(d)

        return b''.join(buf)

    def _read(self, stream):
        return self.parse(stream.read(self.size * 1), self.size, 1, self.signed, self.cstruct.endian)[0]

    def _read_array(self, stream, count):
        return self.parse(stream.read(self.size * count), self.size, count, self.signed, self.cstruct.endian)

    def _read_0(self, stream):
        r = []
        while True:
            v = self._read(stream)
            if v == 0:
                break
            r.append(v)

        return r

    def _write(self, stream, data):
        return stream.write(self.pack([data], self.size, self.cstruct.endian))

    def _write_array(self, stream, data):
        return stream.write(self.pack(data, self.size, self.cstruct.endian))

    def _write_0(self, stream, data):
        return self._write_array(stream, data + [0])

    def default(self):
        return 0

    def default_array(self, count):
        return [0] * count


class Enum(RawType):
    """Implements an Enum type.

    Enums can be made using any type. The API for accessing enums and their
    values is very similar to Python 3 native enums.

    Example:
        When using the default C-style parser, the following syntax is supported:

            enum <name> [: <type>] {
                <values>
            };

        For example, an enum that has A=1, B=5 and C=6 could be written like so:

            enum Test : uint16 {
                A, B=5, C
            };
    """

    def __init__(self, cstruct, name, type_, values):
        self.type = type_
        self.values = values
        self.reverse = {}

        for k, v in values.items():
            self.reverse[v] = k

        super(Enum, self).__init__(cstruct, name, len(self.type))

    def __call__(self, value):
        return EnumInstance(self, value)

    def _read(self, stream):
        v = self.type._read(stream)
        return self(v)

    def _read_array(self, stream, count):
        return list(map(self, self.type._read_array(stream, count)))

    def _read_0(self, stream):
        return list(map(self, self.type._read_0(stream)))

    def _write(self, stream, data):
        data = data.value if isinstance(data, EnumInstance) else data
        return self.type._write(stream, data)

    def _write_array(self, stream, data):
        data = [d.value if isinstance(d, EnumInstance) else d for d in data]
        return self.type._write_array(stream, data)

    def _write_0(self, stream, data):
        data = [d.value if isinstance(d, EnumInstance) else d for d in data]
        return self.type._write_0(stream, data)

    def default(self):
        return self(0)

    def __getitem__(self, attr):
        if attr in self.values:
            return self(self.values[attr])

        raise KeyError(attr)

    def __getattr__(self, attr):
        if attr in self.values:
            return self(self.values[attr])

        raise AttributeError(attr)

    def __contains__(self, attr):
        return attr in self.values


class EnumInstance(object):
    """Implements a value instance of an Enum"""

    def __init__(self, enum, value):
        self.enum = enum
        self.value = value

    @property
    def name(self):
        if self.value not in self.enum.reverse:
            return '{}_{}'.format(self.enum.name, self.value)
        return self.enum.reverse[self.value]

    def __eq__(self, value):
        if isinstance(value, EnumInstance) and value.enum is not self.enum:
            return False

        if hasattr(value, 'value'):
            value = value.value

        return self.value == value

    def __ne__(self, value):
        return self.__eq__(value) is False

    def __hash__(self):
        return hash((self.enum, self.value))

    def __str__(self):
        return '{}.{}'.format(self.enum.name, self.name)

    def __repr__(self):
        return '<{}.{}: {}>'.format(self.enum.name, self.name, self.value)


class Union(RawType):
    def __init__(self, cstruct):
        self.cstruct = cstruct
        super(Union, self).__init__(cstruct)

    def _read(self, stream):
        raise NotImplementedError()


class Pointer(RawType):
    """Implements a pointer to some other type."""

    def __init__(self, cstruct, target):
        self.cstruct = cstruct
        self.type = target
        super(Pointer, self).__init__(cstruct)

    def _read(self, stream, ctx):
        addr = self.cstruct.pointer(stream)
        return PointerInstance(self.type, stream, addr, ctx)

    def __len__(self):
        return len(self.cstruct.pointer)

    def __repr__(self):
        return '<Pointer {!r}>'.format(self.type)


class VoidType(RawType):
    """Implements a void type."""

    def __init__(self):
        super(VoidType, self).__init__(None, 'void', 0)

    def _read(self, stream):
        return None


def ctypes(structure):
    """Create ctypes structures from cstruct structures."""
    fields = []
    for field in structure.fields:
        t = ctypes_type(field.type)
        fields.append((field.name, t))

    tt = type(structure.name, (_ctypes.Structure, ), {"_fields_": fields})
    return tt


def ctypes_type(t):
    mapping = {
        "I": _ctypes.c_ulong,
        "i": _ctypes.c_long,
        "b": _ctypes.c_int8,
    }

    if isinstance(t, PackedType):
        return mapping[t.packchar]

    if isinstance(t, CharType):
        return _ctypes.c_char

    if isinstance(t, Array):
        subtype = ctypes_type(t._type)
        return subtype * t.count

    if isinstance(t, Pointer):
        subtype = ctypes_type(t._target)
        return ctypes.POINTER(subtype)

    raise NotImplementedError("Type not implemented: %s" % t.__class__.__name__)


class Compiler(object):
    """Compiler for cstruct structures. Creates somewhat optimized parsing code."""

    def __init__(self, cstruct):
        self.cstruct = cstruct

    def compile(self, structure):
        source = self.gen_struct_class(structure)
        c = compile(source, '<compiled>', 'exec')

        env = {
            'OrderedDict': OrderedDict,
            'Structure': Structure,
            'Instance': Instance,
            'Expression': Expression,
            'EnumInstance': EnumInstance,
            'PointerInstance': PointerInstance,
            'BytesInteger': BytesInteger,
            'BitBuffer': BitBuffer,
            'struct': struct,
            'xrange': xrange,
        }

        exec(c, env)
        sc = env[structure.name](self.cstruct, structure, source)

        return sc

    def gen_struct_class(self, structure):
        blocks = []
        classes = []
        cur_block = []
        read_size = 0
        prev_was_bits = False

        for field in structure.fields:
            ft = self.cstruct.resolve(field.type)

            if not isinstance(
                ft,
                (
                    Structure,
                    Pointer,
                    Enum,
                    Array,
                    PackedType,
                    CharType,
                    WcharType,
                    BytesInteger,
                ),
            ):
                raise CompilerError("Unsupported type for compiler: {}".format(ft))

            if isinstance(ft, Structure) or (
                isinstance(ft, Array) and isinstance(ft.type, Structure)
            ):
                if cur_block:
                    blocks.append(self.gen_read_block(read_size, cur_block))

                struct_read = 's = stream.tell()\n'
                if isinstance(ft, Array):
                    num = ft.count

                    if isinstance(num, Expression):
                        num = 'max(0, Expression(self.cstruct, "{expr}").evaluate(r))'.format(
                            expr=num.expr
                        )

                    struct_read += (
                        'r["{name}"] = []\n'
                        'for _ in xrange({num}):\n'
                        '    r["{name}"].append(self.cstruct.{struct_name}._read(stream))\n'.format(
                            name=field.name, num=num, struct_name=ft.type.name
                        )
                    )
                else:
                    struct_read += 'r["{name}"] = self.cstruct.{struct_name}._read(stream)\n'.format(
                        name=field.name, struct_name=ft.name
                    )

                struct_read += 'sizes["{name}"] = stream.tell() - s'.format(
                    name=field.name
                )
                blocks.append(struct_read)
                read_size = 0
                cur_block = []
                continue

            if field.bits:
                if cur_block:
                    blocks.append(self.gen_read_block(read_size, cur_block))

                blocks.append(
                    'r["{name}"] = bitreader.read(self.cstruct.{type_name}, {bits})'.format(
                        name=field.name, type_name=field.type.name, bits=field.bits
                    )
                )
                read_size = 0
                cur_block = []
                prev_was_bits = True
                continue
            elif prev_was_bits:
                blocks.append('bitreader.reset()')
                prev_was_bits = False

            try:
                count = len(ft)
                read_size += count
                cur_block.append(field)
            except Exception:
                if cur_block:
                    blocks.append(self.gen_read_block(read_size, cur_block))
                blocks.append(self.gen_dynamic_block(field))
                read_size = 0
                cur_block = []

        if len(cur_block):
            blocks.append(self.gen_read_block(read_size, cur_block))

        read_code = '\n\n'.join(blocks)
        read_code = '\n'.join(['    ' * 2 + line for line in read_code.split('\n')])

        classes.append(COMPILE_TEMPL.format(name=structure.name, read_code=read_code))
        return '\n\n'.join(classes)

    def gen_read_block(self, size, block):
        templ = (
            'buf = stream.read({size})\n'
            'if len(buf) != {size}: raise EOFError()\n'
            'data = struct.unpack(self.cstruct.endian + "{{}}", buf)\n'
            '{{}}'.format(size=size)
        )
        readcode = []
        fmt = []

        curtype = None
        curcount = 0

        buf_offset = 0
        data_offset = 0

        for field in block:
            ft = self.cstruct.resolve(field.type)
            t = ft
            count = 1
            data_count = 1
            read_slice = ''

            if isinstance(t, Enum):
                t = t.type
            elif isinstance(t, Pointer):
                t = self.cstruct.pointer

            if isinstance(ft, Array):
                count = t.count
                data_count = count
                t = t.type

                if isinstance(t, Enum):
                    t = t.type
                elif isinstance(t, Pointer):
                    t = self.cstruct.pointer

                if isinstance(t, (CharType, WcharType, BytesInteger)):
                    read_slice = '{}:{}'.format(
                        buf_offset, buf_offset + (count * t.size)
                    )
                else:
                    read_slice = '{}:{}'.format(data_offset, data_offset + count)
            elif isinstance(t, CharType):
                read_slice = str(buf_offset)
            elif isinstance(t, (WcharType, BytesInteger)):
                read_slice = '{}:{}'.format(buf_offset, buf_offset + t.size)
            else:
                read_slice = str(data_offset)

            if not curtype:
                if isinstance(t, PackedType):
                    curtype = t.packchar
                else:
                    curtype = 'x'

            if isinstance(t, (PackedType, CharType, WcharType, BytesInteger, Enum)):
                charcount = count

                if isinstance(t, (CharType, WcharType, BytesInteger)):
                    data_count = 0
                    packchar = 'x'
                    charcount *= t.size
                else:
                    packchar = t.packchar

                if curtype != packchar:
                    fmt.append('{}{}'.format(curcount, curtype))
                    curcount = 0

                curcount += charcount
                curtype = packchar

            getter = ''
            if isinstance(t, BytesInteger):
                getter = 'BytesInteger.parse(buf[{slice}], {size}, {count}, {signed}, self.cstruct.endian){data_slice}'.format(
                    slice=read_slice,
                    size=t.size,
                    count=count,
                    signed=t.signed,
                    data_slice='[0]' if count == 1 else '',
                )
            elif isinstance(t, (CharType, WcharType)):
                getter = 'buf[{}]'.format(read_slice)
                if isinstance(t, WcharType):
                    getter += ".decode('utf-16-le' if self.cstruct.endian == '<' else 'utf-16-be')"
            else:
                getter = 'data[{}]'.format(read_slice)

            if isinstance(ft, Enum):
                getter = 'EnumInstance(self.cstruct.{type_name}, {getter})'.format(
                    type_name=ft.name, getter=getter
                )
            elif isinstance(ft, Array) and isinstance(ft.type, Enum):
                getter = '[EnumInstance(self.cstruct.{type_name}, d) for d in {getter}]'.format(
                    type_name=ft.type.name, getter=getter
                )
            elif isinstance(ft, Pointer):
                getter = 'PointerInstance(self.cstruct.{type_name}, stream, {getter}, r)'.format(
                    type_name=ft.type.name, getter=getter
                )
            elif isinstance(ft, Array) and isinstance(ft.type, Pointer):
                getter = '[PointerInstance(self.cstruct.{type_name}, stream, d, r) for d in {getter}]'.format(
                    type_name=ft.type.name, getter=getter
                )
            elif isinstance(ft, Array) and isinstance(t, PackedType):
                getter = 'list({})'.format(getter)

            readcode.append(
                'r["{name}"] = {getter}'.format(name=field.name, getter=getter)
            )
            readcode.append(
                'sizes["{name}"] = {size}'.format(name=field.name, size=count * t.size)
            )

            data_offset += data_count
            buf_offset += count * t.size

        if curcount:
            fmt.append('{}{}'.format(curcount, curtype))

        return templ.format(''.join(fmt), '\n'.join(readcode))

    def gen_dynamic_block(self, field):
        if not isinstance(field.type, Array):
            raise CompilerError(
                "Only Array can be dynamic, got {!r}".format(field.type)
            )

        t = field.type.type
        reader = None

        if not field.type.count:  # Null terminated
            if isinstance(t, PackedType):
                reader = (
                    't = []\nwhile True:\n'
                    '    d = stream.read({size})\n'
                    '    if len(d) != {size}: raise EOFError()\n'
                    '    v = struct.unpack(self.cstruct.endian + "{packchar}", d)[0]\n'
                    '    if v == 0: break\n'
                    '    t.append(v)'.format(size=t.size, packchar=t.packchar)
                )

            elif isinstance(t, (CharType, WcharType)):
                reader = (
                    't = []\n'
                    'while True:\n'
                    '    c = stream.read({size})\n'
                    '    if len(c) != {size}: raise EOFError()\n'
                    '    if c == b"{null}": break\n'
                    '    t.append(c)\nt = b"".join(t)'.format(
                        size=t.size, null='\\x00' * t.size
                    )
                )

                if isinstance(t, WcharType):
                    reader += ".decode('utf-16-le' if self.cstruct.endian == '<' else 'utf-16-be')"
            elif isinstance(t, BytesInteger):
                reader = (
                    't = []\n'
                    'while True:\n'
                    '    d = stream.read({size})\n'
                    '    if len(d) != {size}: raise EOFError()\n'
                    '    v = BytesInteger.parse(d, {size}, 1, {signed}, self.cstruct.endian)\n'
                    '    if v == 0: break\n'
                    '    t.append(v)'.format(size=t.size, signed=t.signed)
                )

            return '{reader}\nr["{name}"] = t\nsizes["{name}"] = len(t)'.format(
                reader=reader, name=field.name
            )
        else:
            expr = field.type.count.expr
            expr_read = (
                'dynsize = max(0, Expression(self.cstruct, "{expr}").evaluate(r))\n'
                'buf = stream.read(dynsize * {type_size})\n'
                'if len(buf) != dynsize * {type_size}: raise EOFError()\n'
                'r["{name}"] = {{reader}}\n'
                'sizes["{name}"] = dynsize * {type_size}'.format(
                    expr=expr, name=field.name, type_size=t.size
                )
            )

            if isinstance(t, PackedType):
                reader = 'list(struct.unpack(self.cstruct.endian + "{{:d}}{packchar}".format(dynsize), buf))'.format(
                    packchar=t.packchar, type_size=t.size
                )
            elif isinstance(t, (CharType, WcharType)):
                reader = 'buf'
                if isinstance(t, WcharType):
                    reader += ".decode('utf-16-le' if self.cstruct.endian == '<' else 'utf-16-be')"
            elif isinstance(t, BytesInteger):
                reader = 'BytesInteger.parse(buf, {size}, dynsize, {signed}, self.cstruct.endian)'.format(
                    size=t.size, signed=t.signed
                )

            return expr_read.format(reader=reader, size=None)


def hexdump(s, palette=None, offset=0, prefix=""):
    """Hexdump some data.

    Args:
        s: Bytes to hexdump.
        palette: Colorize the hexdump using this color pattern.
        offset: Byte offset of the hexdump.
        prefix: Optional prefix.
    """
    if palette:
        palette = palette[::-1]

    remaining = 0
    active = None

    for i in xrange(0, len(s), 16):
        vals = ""
        chars = []
        for j in xrange(16):
            if not active and palette:
                remaining, active = palette.pop()
                vals += active
            elif active and j == 0:
                vals += active

            if i + j >= len(s):
                vals += "  "
            else:
                c = s[i + j]
                c = chr(c) if PY3 else c
                p = c if c in PRINTABLE else "."

                if active:
                    vals += "{:02x}".format(ord(c))
                    chars.append(active + p + COLOR_NORMAL)
                else:
                    vals += "{:02x}".format(ord(c))
                    chars.append(p)

                remaining -= 1
                if remaining == 0:
                    active = None

                    if palette is not None:
                        vals += COLOR_NORMAL

                if j == 15:
                    if palette is not None:
                        vals += COLOR_NORMAL

            vals += " "

            if j == 7:
                vals += " "

        chars = "".join(chars)
        print("{}{:08x}  {:48s}  {}".format(prefix, offset + i, vals, chars))


def dumpstruct(t, data=None, offset=0):
    """Dump a structure or parsed structure instance.

    Prints a colorized hexdump and parsed structure output.

    Args:
        t: Structure or Instance to dump.
        data: Bytes to parse the Structure on, if t is not a parsed Instance.
        offset: Byte offset of the hexdump.
    """
    colors = [
        (COLOR_RED, COLOR_BG_RED),
        (COLOR_GREEN, COLOR_BG_GREEN),
        (COLOR_YELLOW, COLOR_BG_YELLOW),
        (COLOR_BLUE, COLOR_BG_BLUE),
        (COLOR_PURPLE, COLOR_BG_PURPLE),
        (COLOR_CYAN, COLOR_BG_CYAN),
        (COLOR_WHITE, COLOR_BG_WHITE),
    ]

    if isinstance(t, Instance):
        g = t
        t = t._type
        data = g.dumps()
    elif isinstance(t, Structure) and data:
        g = t(data)
    else:
        raise ValueError("Invalid arguments")

    palette = []
    ci = 0
    out = "struct {}".format(t.name) + ":\n"
    for field in g._type.fields:
        fg, bg = colors[ci % len(colors)]
        palette.append((g._size(field.name), bg))
        ci += 1

        v = getattr(g, field.name)
        if isinstance(v, str):
            v = repr(v)
        elif isinstance(v, int):
            v = hex(v)
        elif isinstance(v, list):
            v = pprint.pformat(v)
            if '\n' in v:
                v = v.replace('\n', '\n{}'.format(' ' * (len(field.name) + 4)))

        out += "- {}{}{}: {}\n".format(fg, field.name, COLOR_NORMAL, v)

    print()
    hexdump(data, palette, offset=offset)
    print()
    print(out)