from filetypes.base import *
import malcat
import struct



class SyskindRecord(Struct):

    def parse(self):
        id = yield UInt16(name="Id", comment="Record ID")
        if id != 1:
            raise FatalError
        sz = yield UInt32(name="Size", comment="Record size")
        if sz != 4:
            raise FatalError
        yield UInt32(name="Syskind", comment="Project Syskind", values=[
            ("16-bits Windows", 0),
            ("32-bits Windows", 1),
            ("Macintosh", 2),
            ("64-bits Windows", 3),
        ])

class CompatVersionRecord(Struct):

    def parse(self):
        id = yield UInt16(name="Id", comment="Record ID")
        if id != 0x4a:
            raise FatalError
        sz = yield UInt32(name="Size", comment="Record size")
        if sz != 4:
            raise FatalError
        yield UInt32(name="CompatVersion", comment="specifies the compat version value for the VBA project")        

class LcidRecord(Struct):

    def parse(self):
        id = yield UInt16(name="Id", comment="Record ID")
        if id != 2:
            raise FatalError
        sz = yield UInt32(name="Size", comment="Record size")
        if sz != 4:
            raise FatalError
        yield UInt32(name="Lcid", comment="Project LCID")        


class LcidInvokeRecord(Struct):

    def parse(self):
        id = yield UInt16(name="Id", comment="Record ID")
        if id != 0x14:
            raise FatalError
        sz = yield UInt32(name="Size", comment="Record size")
        if sz != 4:
            raise FatalError
        yield UInt32(name="LcidInvoke", comment="Project invoke LCID")         


class CodePageRecord(Struct):

    def parse(self):
        id = yield UInt16(name="Id", comment="Record ID")
        if id != 3:
            raise FatalError
        sz = yield UInt32(name="Size", comment="Record size")
        if sz != 2:
            raise FatalError
        yield UInt16(name="CodePage", comment="Project codepage")            


class NameRecord(Struct):

    def __init__(self, id, *args, **kwargs):
        self.id = id
        Struct.__init__(self, *args, **kwargs)

    def parse(self):
        id = yield UInt16(name="Id", comment="Record ID")
        if id != self.id:
            raise FatalError
        sz = yield UInt32(name="Size", comment="Record size")
        yield String(sz, name="Name", comment="Name")


class UnicodeNameRecord(Struct):

    def __init__(self, id, *args, **kwargs):
        self.id = id
        Struct.__init__(self, *args, **kwargs)

    def parse(self):
        id = yield UInt16(name="Id", comment="Record ID")
        if id != self.id:
            raise FatalError
        sz = yield UInt32(name="Size", comment="Record size")
        yield StringUtf16le(sz // 2, name="Name", comment="Name")        


class DocStringRecord(Struct):

    def parse(self):
        id = yield UInt16(name="Id", comment="Record ID")
        sz = yield UInt32(name="SizeOfDocString", comment="size of ascii doc string")
        if sz:
            yield String(sz, name="DocString", comment="docstring")        
        reserved = yield UInt16(name="Reserved")
        if reserved != 0x40 and reserved != 0x48:
            raise FatalError
        sz = yield UInt32(name="SizeOfUnicodeDocString", comment="size of unicode doc string")
        if sz:
            yield StringUtf16le(sz // 2, name="UnicodeDocString", comment="unicode docstring")


class HelpFileRecord(Struct):

    def parse(self):
        id = yield UInt16(name="Id", comment="Record ID")
        if id != 6:
            raise FatalError
        sz = yield UInt32(name="SizeOfHelpFile1", comment="size of first help file path")
        if sz:
            yield String(sz, name="HelpFile1", comment="Project first help file path")
        reserved = yield UInt16(name="Reserved")
        if reserved != 0x3d:
            raise FatalError
        sz = yield UInt32(name="SizeOfHelpFile2", comment="size of second help file path")
        if sz:
            yield String(sz, name="HelpFile2", comment="Project second help file path")        

class HelpContextRecord(Struct):

    def parse(self):
        id = yield UInt16(name="Id", comment="Record ID")
        if id != 7 and id != 0x1e:
            raise FatalError
        sz = yield UInt32(name="Size", comment="Record size")
        if sz != 4:
            raise FatalError
        yield UInt32(name="HelpContext", comment="Unsigned integer that specifies the Help topic identifier in the Help file")         


class LibFlagsRecord(Struct):

    def parse(self):
        id = yield UInt16(name="Id", comment="Record ID")
        if id != 8:
            raise FatalError
        sz = yield UInt32(name="Size", comment="Record size")
        if sz != 4:
            raise FatalError
        yield UInt32(name="ProjectLibFlags", comment="An unsigned integer that specifies LIBFLAGS for the VBA project’s Automation type library as specified in [MS-OAUT] section 2.2.20. MUST be 0x00000000.")        


class VersionRecord(Struct):

    def parse(self):
        id = yield UInt16(name="Id", comment="Record ID")
        if id != 9:
            raise FatalError
        sz = yield UInt32(name="Reserved")
        if sz != 4:
            raise FatalError
        yield UInt32(name="VersionMajor")
        yield UInt16(name="VersionMinor")


class ConstantRecord(Struct):

    def parse(self):
        id = yield UInt16(name="Id", comment="Record ID")
        if id != 12:
            raise FatalError
        sz = yield UInt32(name="SizeOfConstants")
        if sz:
            yield String(sz, name="Constants")
        reserved = yield UInt16(name="Reserved")
        if reserved != 0x3c:
            raise FatalError
        sz = yield UInt32(name="SizeOfConstantsUnicode")
        if sz:
            yield StringUtf16le(sz // 2, name="ConstantsUnicode")




class ProjectInformation(Struct):

    def parse(self):
        yield SyskindRecord()
        if self.look_ahead(2) == b"\x4a\x00":
            yield CompatVersionRecord()
        yield LcidRecord()
        yield LcidInvokeRecord()
        yield CodePageRecord()
        yield NameRecord(4)
        yield DocStringRecord()
        yield HelpFileRecord()
        yield HelpContextRecord()
        yield LibFlagsRecord()
        yield VersionRecord()
        yield ConstantRecord()            


class ReferenceNameRecord(Struct):

    def parse(self):
        id = yield UInt16(name="Id", comment="Record ID")
        sz = yield UInt32(name="SizeOfName")
        if sz:
            yield String(sz, name="Name")
        reserved = yield UInt16(name="Reserved")
        if reserved != 0x3e:
            raise FatalError
        sz = yield UInt32(name="SizeOfNameUnicode")
        if sz:
            yield StringUtf16le(sz // 2, name="NameUnicode")


class ReferenceRegisteredRecord(Struct):

    def parse(self):
        id = yield UInt16(name="Id", comment="Record ID")
        yield UInt32(name="Size", comment="An unsigned integer that specifies the total size in bytes of SizeOfLibid, Libid, Reserved1, and Reserved2. MUST be ignored on read")
        sz = yield UInt32(name="SizeOfLibid")
        yield String(sz, name="Libid", comment="An array of SizeOfLibid bytes that specifies an Automation type library’s identifier. MUST contain MBCS characters encoded using the code page specified in PROJECTCODEPAGE")
        yield Unused(6)


class ReferenceOriginalRecord(Struct):

    def parse(self):
        id = yield UInt16(name="Id", comment="Record ID")
        sz = yield UInt32(name="SizeOfLibid")
        yield String(sz, name="Libid", comment="specifies the identifier of the Automation type library a REFERENCECONTROL (section 2.3.4.2.2.3) was generated from")


class ReferenceProjectRecord(Struct):

    def parse(self):
        id = yield UInt16(name="Id", comment="Record ID")
        yield UInt32(name="Size", comment="An unsigned integer that specifies the total size in bytes of Record. MUST be ignored on read")
        sz = yield UInt32(name="SizeOfLibidAbsolute")
        if sz:
            yield String(sz, name="LibidAbsolute", comment="The referenced VBA project’s identifier with an absolute path, <ProjectPath>")
        sz = yield UInt32(name="SizeOfLibidRelative")
        if sz:
            yield String(sz, name="LibidRelative", comment="The eferenced VBA project’s identifier with a relative path, <ProjectPath>, that is relative to the current VBA project")
        yield UInt32(name="VersionMajor", comment="Major version of referenced VBA project")
        yield UInt16(name="VersionMinor", comment="Minor version of referenced VBA project")


class ReferenceControlRecord(Struct):

    def parse(self):
        id = yield UInt16(name="Id", comment="Record ID")
        yield UInt32(name="SizeTwiddled ", comment="An unsigned integer that specifies the total size in bytes of Record. MUST be ignored on read")
        sz = yield UInt32(name="SizeOfLibidTwiddled")
        if sz:
            yield String(sz, name="LibidTwiddled", comment="SHOULD be '*\G{00000000-0000-0000-0000-000000000000}#0.0#0##' (case-sensitive). MAY<8> specify a twiddled type library’s identifier.")
        yield Unused(6)
        data, = struct.unpack("<H", self.look_ahead(2))
        if data == 0x16:
            yield ReferenceNameRecord(name="NameRecordExtended", comment="Specifies the name of the extended type library")
        reserved = yield UInt16(name="Reserved")
        if reserved != 0x30:
            raise FatalError
        yield UInt32(name="SizeExtended")
        sz = yield UInt32(name="SizeOfLibidExtended")
        if sz:
            yield String(sz, name="LibidExtended", comment="Specifies the extended type library’s identifier")
        yield Unused(6)
        yield GUID(name="OriginalTypeLib", comment="Specifies the Automation type library the extended type library was generated from")
        yield UInt32(name="Cookie", comment="An unsigned integer that specifies the extended type library’s cookie. MUST be unique for each REFERENCECONTROL")


class References(Struct):

    def parse(self):
        REFID_TO_TYPE = {
            0x16: ReferenceNameRecord,
            0xe: ReferenceProjectRecord,
            0x2f: ReferenceControlRecord,
            0x33: ReferenceOriginalRecord,
            0xd: ReferenceRegisteredRecord
        }
        while self.remaining() > 2:
            recordid, = struct.unpack("<H", self.look_ahead(2))
            if recordid == 0xf:
                break
            class_ = REFID_TO_TYPE.get(recordid, None)
            if class_ is None:
                raise FatalError("Unrecognized record id: {:x}".format(recordid))
            yield class_()


class CookieRecord(Struct):

    def parse(self):
        id = yield UInt16(name="Id", comment="Record ID")
        sz = yield UInt32(name="Size")
        if sz != 2:
            raise FatalError
        yield UInt16(name="Cookie")


class ModuleOffsetRecord(Struct):

    def parse(self):
        id = yield UInt16(name="Id", comment="Record ID")
        if id != 0x31:
            raise FatalError
        sz = yield UInt32(name="Size", comment="Record size")
        if sz != 4:
            raise FatalError
        yield UInt32(name="TextOffset", comment="An unsigned integer that specifies the byte offset of the source code in the ModuleStream")         

class AtomicRecord(Struct):

    def parse(self):
        id = yield UInt16(name="Id", comment="Record ID")
        yield UInt32(name="Reserved")


class ModuleRecord(Struct):

    def parse(self):
        yield NameRecord(0x19, name="ModuleNameRecord")
        recordid, = struct.unpack("<H", self.look_ahead(2))
        if recordid == 0x47:
            yield UnicodeNameRecord(0x47, name="ModuleUnicodeNameRecord")
        yield NameRecord(0x1a, name="ModuleStreamNameRecord")
        recordid, = struct.unpack("<H", self.look_ahead(2))
        if recordid == 0x32:
            yield UnicodeNameRecord(0x32, name="ModuleStreamNameUnicodeRecord")
        yield DocStringRecord()
        yield ModuleOffsetRecord()
        yield HelpContextRecord()
        yield CookieRecord()
        yield AtomicRecord(name="ModuleTypeRecord")
        recordid, = struct.unpack("<H", self.look_ahead(2))
        if recordid == 0x25:
            yield AtomicRecord(name="ReadOnlyRecord")
        recordid, = struct.unpack("<H", self.look_ahead(2))
        if recordid == 0x28:
            yield AtomicRecord(name="PrivateRecord")
        eob = yield UInt16(name="Terminator")
        if eob != 0x2b:
            raise FatalError
        yield Unused(4, name="Reserved")


class Modules(Struct):

    def parse(self):
        id = yield UInt16(name="Id", comment="Record ID")
        sz = yield UInt32(name="Size")
        if sz != 2:
            raise FatalError
        num_modules = yield UInt16(name="Count", comment="Number of modules")
        yield CookieRecord()
        for i in range(num_modules):
            yield ModuleRecord(name="Module[{}]".format(i))


class VBADirAnalyzer(FileTypeAnalyzer):
    category = malcat.FileType.DOCUMENT
    name = "Office.VbaDir"
    regexp = r"\x01\x00\x04\x00\x00\x00.\x00\x00\x00\x02\x00\x04\x00\x00\x00....\x14\x00\x04\x00\x00\x00....\x03\x00\x02\x00\x00\x00" 


    def parse(self, hint):
        self.confirm()
        yield ProjectInformation()
        yield References()
        yield Modules()
        eob = yield UInt16(name="Terminator")
        if eob != 0x10:
            raise FatalError
        yield Unused(4, name="Reserved")


    



