notded/dotnet/file.py

import pefile
import destruct
from destruct import Type, Struct, Arr

from .common import CLRSectionReference, CLRToken, CLRStreamType, CLRStreamIndex, CLRTableIndex, CLRTableRange
from .streams import STREAM_PARSERS


class CLRHeader(Struct):
    size = UInt(32)
    version_major = UInt(16)
    version_minor = UInt(16)
    metadata_info = CLRSectionReference
    flags = UInt(32)
    entrypoint = CLRToken
    unk_info = CLRSectionReference
    namesig_info = CLRSectionReference
    _gap0x28 = Data(32)

class CLRStreamMetadata(Struct):
    offset = UInt(32)
    size   = UInt(32)
    name   = AlignTo(Str(kind='c'), 4)

class CLRMetadataHeader(Struct):
    magic = Sig(b'BSJB')
    version_major = UInt(16)
    version_minor = UInt(16)
    _gap0x8 = Data(4)
    version = Str(kind='pascal', length_type=UInt(32))
    _gap0x12 = Data(2)
    stream_count = UInt(16)
    streams = Arr(CLRStreamMetadata)

    def on_stream_count(self, spec, context):
        spec.streams.count = self.stream_count

class CLRTableWrapper:
    __slots__ = ('__file', '__table', '__index', '__child')

    def __init__(self, file, table, index, child):
        self.__file = file
        self.__table = table
        self.__index = index
        self.__child = child

    def __getattr__(self, name):
        val = getattr(self.__child, name)
        if isinstance(self.__child, destruct.Struct):
            type = self.__child._spec[name]
            if isinstance(type, CLRStreamIndex):
                stream = type.type
                if stream == CLRStreamType.String:
                    val = self.__file.get_string_at(val)
                elif stream == CLRStreamType.Blob:
                    val = self.__file.get_blob_at(val)
                elif stream == CLRStreamType.UserString:
                    val = self.__file.get_user_string_at(val)
                elif stream == CLRStreamType.GUID:
                    val = self.__file.get_guid_at(val)
                if type.child:
                    c = destruct.Context(type.child)
                    c.user.metadata = self.__file.streams[CLRStreamType.Metadata]
                    val = destruct.parse(type.child, val, c)
                return val
            elif isinstance(type, CLRTableIndex):
                table = type.type
                return self.__file.get_table_entry(table, val)
            elif isinstance(type, CLRTableRange):
                table = type.type
                if self.__index + 1 < self.__file.get_table_size(self.__table):
                    next = self.__file.get_table_entry(self.__table, self.__index + 1)
                    end = min(self.__file.get_table_size(table), getattr(next.__child, name))
                else:
                    end = self.__file.get_table_size(table)
                return [self.__file.get_table_entry(table, i) for i in range(val, end)]
        return val

    def __repr__(self):
        return '<{}: {!r} in {!r}>'.format(self.__class__.__name__, self.__child, self.__file)

class CLRFile:
    def __init__(self, fn):
        self.name = fn
        self.pe = pefile.PE(fn)
        self.header = None
        self.metadata = None
        self.streams = {}
        self.parse()

    def __repr__(self):
        return '<{}: "{}">'.format(self.__class__.__name__, self.name)

    def parse_at(self, c, offset, size=None):
        if size is None:
            size = destruct.sizeof(c)
        buf = self.pe.get_data(rva=offset, length=size)
        return destruct.parse(c, buf)

    def parse(self):
        net_data_entry = pefile.DIRECTORY_ENTRY['IMAGE_DIRECTORY_ENTRY_COM_DESCRIPTOR']
        if net_data_entry < len(self.pe.OPTIONAL_HEADER.DATA_DIRECTORY):
            net_data_dir = self.pe.OPTIONAL_HEADER.DATA_DIRECTORY[net_data_entry]
            net_offset = net_data_dir.VirtualAddress
            net_size = net_data_dir.Size
        else:
            for section in self.pe.sections:
                if section.Name.rstrip(b'\x00') != b'.text':
                    continue
                net_offset = section.PointerToRawData + 8
                net_size = 72

        self.header = self.parse_at(CLRHeader, net_offset, net_size)
        self.metadata = self.parse_at(CLRMetadataHeader, self.header.metadata_info.rva, self.header.metadata_info.size)
        self.streams = {}
        for s in self.metadata.streams:
            type = CLRStreamType(s.name)
            if type in STREAM_PARSERS:
                stream = self.parse_at(STREAM_PARSERS[type], self.header.metadata_info.rva + s.offset, s.size)
            else:
                stream = None
            self.streams[type] = stream

    def get_table(self, t):
        return (CLRTableWrapper(self, t, i, x) for i, x in enumerate(self.streams[CLRStreamType.Metadata].tables[t]))

    def get_table_size(self, t):
        return len(self.streams[CLRStreamType.Metadata].tables[t])

    def get_table_entry(self, t, i):
        return CLRTableWrapper(self, t, i, self.streams[CLRStreamType.Metadata].tables[t][i])

    def get_string_at(self, i):
        buf = bytearray()
        while True:
            c = self.streams[CLRStreamType.String].data[i]
            if not c:
                break
            buf.append(c)
            i += 1
        return buf.decode('utf-8')

    def _get_length_prefixed_value(self, stream, i):
        length = self.streams[stream].data[i]
        if (length >> 5) == 0b110:
            length = length & 0b11111
            nbytes = 4
        elif (length >> 6) == 0b10:
            length = length & 0b111111
            nbytes = 2
        else:
            nbytes = 1
        for off in range(nbytes - 1):
            length = (length << 8) | self.streams[stream].data[i + 1 + off]
        return self.streams[stream].data[i + nbytes:i + nbytes + length]

    def get_blob_at(self, i):
        return self._get_length_prefixed_value(CLRStreamType.Blob, i)

    def get_user_string_at(self, i):
        return self._get_length_prefixed_value(CLRStreamType.UserString, i)[:-1].decode('utf-16le')

    def get_guid_at(self, i):
        if i == 0:
            return None
        return self.streams[CLRStreamType.GUID].guids[i - 1]

    def get_by_token(self, t):
        if not t.row:
            return None
        return self.get_table_entry(t.table, t.row - 1)

    @property
    def entrypoint(self):
        return self.get_by_token(self.header.entrypoint)