Ticket #138: ole2_util.py

from hachoir_core.endian import LITTLE_ENDIAN
from hachoir_core.field import Int8, RawBytes, RootSeekableFieldSet, FieldSet, ParserError
from hachoir_core.stream import StringInputStream
from hachoir_parser import HachoirParser

class OLE2FragmentParser(HachoirParser,RootSeekableFieldSet):
    tags = {
        "description": "Microsoft Office document subfragments",
    }
    endian = LITTLE_ENDIAN
    

    OS_MAC = 1
    OS_NAME = {
        0: "Windows 16-bit",
        1: "Macintosh",
        2: "Windows 32-bit",
    }
    ENDIAN_CHECK=False
    OS_CHECK=False

    def __init__(self, stream, **args):
        RootSeekableFieldSet.__init__(self, None, "root", stream, None, stream.askSize(self))
        HachoirParser.__init__(self, stream, **args)
        if self.ENDIAN_CHECK:
            if self["endian"].value == "\xFF\xFE":
                self.endian = BIG_ENDIAN
            elif self["endian"].value == "\xFE\xFF":
                self.endian = LITTLE_ENDIAN
            else:
                raise ParserError("OLE2: Invalid endian value")
        if self.OS_CHECK:
            self.osconfig = OSConfig(self["os"].value == self.OS_MAC)

    def validate(self):
        if self.ENDIAN_CHECK:
            if self["endian"].value not in ["\xFF\xFE", "\xFE\xFF"]:
                return "Unknown endian value %s"%self["endian"].value.encode('hex')
        return True

class RawParser(OLE2FragmentParser):
    ENDIAN_CHECK=False
    OS_CHECK=False
    def createFields(self):
        yield RawBytes(self,"rawdata",self.datasize)
        if self.datasize<self.size//8: yield RawBytes(self,"slack_space",(self.size//8)-self.datasize)

class OSConfig:
    def __init__(self, big_endian):
        if big_endian:
            self.charset = "MacRoman"
            self.utf16 = "UTF-16-BE"
        else:
            self.charset = "ISO-8859-1"
            self.utf16 = "UTF-16-LE"

class Bool(Int8):
    def createValue(self):
        value = Int8.createValue(self)
        return (value == -1)

class FragmentGroup:
    def __init__(self, parser):
        self.items = []
        self.parser = parser
        self.args = {}

    def add(self, item):
        self.items.append(item)

    def createInputStream(self):
        # FIXME: Use lazy stream creation
        data = []
        for item in self.items:
            data.append( item["rawdata"].value )
        data = "".join(data)

        # FIXME: Use smarter code to send arguments
        self.args["ole2"] = self.items[0].root
        tags = {"class": self.parser, "args": self.args}
        tags = tags.iteritems()
        return StringInputStream(data, "<fragment group>", tags=tags)

class CustomFragment(FieldSet):
    def __init__(self, parent, name, size, parser, description=None, group=None):
        FieldSet.__init__(self, parent, name, description, size=size)
        if not group:
            group = FragmentGroup(parser)
        self.group = group
        self.group.add(self)

    def createFields(self):
        yield RawBytes(self, "rawdata", self.size//8)

    def _createInputStream(self, **args):
        return self.group.createInputStream()