from filetypes.base import *
import re
import struct

class GOLayoutTable32(Struct):

    def parse(self):
        yield Va32(name="PcLnAddress", comment="address of pcln header")
        yield UInt32(name="PcLnSize", comment="size of pcln table")
        yield UInt32(name="PcLnSize2", comment="size of pcln table")
        yield Va32(name="PcLnFunctionAddress", comment="address of function table (inside pcln table)")
        yield UInt32(name="PcLnFunctionEntries", comment="number of entries in pcln function table")
        yield UInt32(name="PcLnFunctionEntries2", comment="number of entries in pcln function table")
        yield Va32(name="PcLnFileAddress", comment="address of pcln source file table (inside pcln table)")
        yield UInt32(name="PcLnFileEntries", comment="number of entries in pcln file table")
        yield UInt32(name="PcLnFileEntries2", comment="number of entries in pcln file table")


class GOLayoutTable64(Struct):

    def parse(self):
        yield Va64(name="PcLnAddress", comment="address of pcln header")
        yield UInt64(name="PcLnSize", comment="size of pcln struct, including header")
        yield UInt64(name="PcLnSize2", comment="size of pcln struct, including header")
        yield Va64(name="PcLnFunctionAddress", comment="address of function table (inside pcln table)")
        yield UInt64(name="PcLnFunctionEntries", comment="number of entries in pcln function table")
        yield UInt64(name="PcLnFunctionEntries2", comment="number of entries in pcln function table")
        yield Va64(name="PcLnFileAddress", comment="address of pcln source file table (inside pcln table)")
        yield UInt64(name="PcLnFileEntries", comment="number of entries in pcln file table")
        yield UInt64(name="PcLnFileEntries2", comment="number of entries in pcln file table")


class GOLayoutTable32_116(Struct):

    def parse(self):
        yield Va32(name="PcLnAddress", comment="address of pcln header")
        yield Va32(name="PcLnFunctionNameAddress", comment="address of function name table (inside pcln table)")
        yield UInt32(name="PcLnFunctionNameEntries", comment="number of entries in pcln function name table")
        yield UInt32(name="PcLnFunctionNameEntries", comment="number of entries in pcln function name table")
        yield Va32(name="PcLnCuAddress", comment="address of cu table (inside pcln table)")
        yield UInt32(name="PcLnCuEntries", comment="number of entries in pcln cu table")
        yield UInt32(name="PcLnCuEntries2", comment="number of entries in pcln cu table")
        yield Va32(name="PcLnFileAddress", comment="address of pcln source file table (inside pcln table)")
        yield UInt32(name="PcLnFileEntries", comment="number of entries in pcln file table")
        yield UInt32(name="PcLnFileEntries2", comment="number of entries in pcln file table")
        yield Va32(name="PcLnPcAddress", comment="address of pcln source pc table (inside pcln table)")
        yield UInt32(name="PcLnPcEntries", comment="number of entries in pcln pc table")
        yield UInt32(name="PcLnPcEntries2", comment="number of entries in pcln pc table")
        yield Va32(name="PcLnFunctionAddress", comment="address of pcln source function table (inside pcln table)")
        yield UInt32(name="PcLnFunctionEntries", comment="number of entries in function file table")
        yield UInt32(name="PcLnFunctionEntries2", comment="number of entries in function file table")


class GOLayoutTable64_116(Struct):

    def parse(self):
        yield Va64(name="PcLnAddress", comment="address of pcln header")
        yield Va64(name="PcLnFunctionNameAddress", comment="address of function name table (inside pcln table)")
        yield UInt64(name="PcLnFunctionNameEntries", comment="number of entries in pcln function name table")
        yield UInt64(name="PcLnFunctionNameEntries", comment="number of entries in pcln function name table")
        yield Va64(name="PcLnCuAddress", comment="address of cu table (inside pcln table)")
        yield UInt64(name="PcLnCuEntries", comment="number of entries in pcln cu table")
        yield UInt64(name="PcLnCuEntries2", comment="number of entries in pcln cu table")
        yield Va64(name="PcLnFileAddress", comment="address of pcln source file table (inside pcln table)")
        yield UInt64(name="PcLnFileEntries", comment="number of entries in pcln file table")
        yield UInt64(name="PcLnFileEntries2", comment="number of entries in pcln file table")
        yield Va64(name="PcLnPcAddress", comment="address of pcln source pc table (inside pcln table)")
        yield UInt64(name="PcLnPcEntries", comment="number of entries in pcln pc table")
        yield UInt64(name="PcLnPcEntries2", comment="number of entries in pcln pc table")
        yield Va64(name="PcLnFunctionAddress", comment="address of pcln source function table (inside pcln table)")
        yield UInt64(name="PcLnFunctionEntries", comment="number of entries in function file table")
        yield UInt64(name="PcLnFunctionEntries2", comment="number of entries in function file table")


class GOLayoutTable32_118(Struct):

    def parse(self):
        yield Va32(name="PcLnAddress", comment="address of pcln header")
        yield Va32(name="PcLnFunctionNameAddress", comment="address of function name table (inside pcln table)")
        yield UInt32(name="PcLnFunctionNameEntries", comment="number of entries in pcln function name table")
        yield UInt32(name="PcLnFunctionNameEntries", comment="number of entries in pcln function name table")
        yield Va32(name="PcLnCuAddress", comment="address of cu table (inside pcln table)")
        yield UInt32(name="PcLnCuEntries", comment="number of entries in pcln cu table")
        yield UInt32(name="PcLnCuEntries2", comment="number of entries in pcln cu table")
        yield Va32(name="PcLnFileAddress", comment="address of pcln source file table (inside pcln table)")
        yield UInt32(name="PcLnFileEntries", comment="number of entries in pcln file table")
        yield UInt32(name="PcLnFileEntries2", comment="number of entries in pcln file table")
        yield Va32(name="PcLnPcAddress", comment="address of pcln source pc table (inside pcln table)")
        yield UInt32(name="PcLnPcEntries", comment="number of entries in pcln pc table")
        yield UInt32(name="PcLnPcEntries2", comment="number of entries in pcln pc table")
        yield Va32(name="PcLnFunctionAddress", comment="address of pcln source function table (inside pcln table)")
        yield UInt32(name="PcLnFunctionEntries", comment="number of entries in function file table")
        yield UInt32(name="PcLnFunctionEntries2", comment="number of entries in function file table")


class GOLayoutTable64_118(Struct):

    def parse(self):
        yield Va64(name="PcLnAddress", comment="address of pcln header")
        yield Va64(name="PcLnFunctionNameAddress", comment="address of function name table (inside pcln table)")
        yield UInt64(name="PcLnFunctionNameEntries", comment="number of entries in pcln function name table")
        yield UInt64(name="PcLnFunctionNameEntries", comment="number of entries in pcln function name table")
        yield Va64(name="PcLnCuAddress", comment="address of cu table (inside pcln table)")
        yield UInt64(name="PcLnCuEntries", comment="number of entries in pcln cu table")
        yield UInt64(name="PcLnCuEntries2", comment="number of entries in pcln cu table")
        yield Va64(name="PcLnFileAddress", comment="address of pcln source file table (inside pcln table)")
        yield UInt64(name="PcLnFileEntries", comment="number of entries in pcln file table")
        yield UInt64(name="PcLnFileEntries2", comment="number of entries in pcln file table")
        yield Va64(name="PcLnPcAddress", comment="address of pcln source pc table (inside pcln table)")
        yield UInt64(name="PcLnPcEntries", comment="number of entries in pcln pc table")
        yield UInt64(name="PcLnPcEntries2", comment="number of entries in pcln pc table")
        yield Va64(name="PcLnFunctionAddress", comment="address of pcln source function table (inside pcln table)")
        yield UInt64(name="PcLnFunctionEntries", comment="number of entries in function file table")
        yield UInt64(name="PcLnFunctionEntries2", comment="number of entries in function file table")



class GOPCLnHeader(Struct):

    def parse(self):
        yield Bytes(6, name="Signature", comment="always F?FFFFFF0000")
        yield UInt8(name="MinimumOpcodeSize")
        yield UInt8(name="PointerSize", comment="size of a pointer infunction table")
        yield UInt32(name="NumberOfFunctions", comment="number of functions in function table")
        if self.look_ahead(4) == b"\x00\x00\x00\x00":
            yield Unused(4, name="Padding")

class GOPCLnHeader32_116(Struct):

    def parse(self):
        yield Bytes(6, name="Signature", comment="always F?FFFFFF0000")
        yield UInt8(name="MinimumOpcodeSize")
        yield UInt8(name="PointerSize", comment="size of a pointer infunction table")
        yield UInt32(name="NumberOfFunctions", comment="number of functions in function table")
        yield UInt32(name="NumberOfFiles", comment="number of functions in function table")
        yield Offset32(name="FunctionNameTable", base=self.offset, comment="")
        yield Offset32(name="CuTable", base=self.offset, comment="")
        yield Offset32(name="FileTable", base=self.offset, comment="")
        yield Offset32(name="PcTable", base=self.offset, comment="")
        yield Offset32(name="FunctionTable", base=self.offset, comment="")


class GOPCLnHeader64_116(Struct):

    def parse(self):
        yield Bytes(6, name="Signature", comment="always F?FFFFFF0000")
        yield UInt8(name="MinimumOpcodeSize")
        yield UInt8(name="PointerSize", comment="size of a pointer infunction table")
        yield UInt64(name="NumberOfFunctions", comment="number of functions in function table")
        yield UInt64(name="NumberOfFiles", comment="number of functions in function table")
        yield Offset64(name="FunctionNameTable", base=self.offset, comment="")
        yield Offset64(name="CuTable", base=self.offset, comment="")
        yield Offset64(name="FileTable", base=self.offset, comment="")
        yield Offset64(name="PcTable", base=self.offset, comment="")
        yield Offset64(name="FunctionTable", base=self.offset, comment="")

class GOPCLnHeader32_118(Struct):

    def parse(self):
        yield Bytes(6, name="Signature", comment="always F?FFFFFF0000")
        yield UInt8(name="MinimumOpcodeSize")
        yield UInt8(name="PointerSize", comment="size of a pointer infunction table")
        yield UInt32(name="NumberOfFunctions", comment="number of functions in function table")
        yield UInt32(name="NumberOfFiles", comment="number of functions in function table")
        yield Va32(name="TextStart", comment="")
        yield Offset32(name="FunctionNameTable", base=self.offset, comment="")
        yield Offset32(name="CuTable", base=self.offset, comment="")
        yield Offset32(name="FileTable", base=self.offset, comment="")
        yield Offset32(name="PcTable", base=self.offset, comment="")
        yield Offset32(name="FunctionTable", base=self.offset, comment="")


class GOPCLnHeader64_118(Struct):

    def parse(self):
        yield Bytes(6, name="Signature", comment="always F?FFFFFF0000")
        yield UInt8(name="MinimumOpcodeSize")
        yield UInt8(name="PointerSize", comment="size of a pointer infunction table")
        yield UInt64(name="NumberOfFunctions", comment="number of functions in function table")
        yield UInt64(name="NumberOfFiles", comment="number of functions in function table")
        yield Va64(name="TextStart", comment="")
        yield Offset64(name="FunctionNameTable", base=self.offset, comment="")
        yield Offset64(name="CuTable", base=self.offset, comment="")
        yield Offset64(name="FileTable", base=self.offset, comment="")
        yield Offset64(name="PcTable", base=self.offset, comment="")
        yield Offset64(name="FunctionTable", base=self.offset, comment="")        
        

class GOPCLnFunctionInfoPointer32(Struct):

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

    def parse(self):
        yield Va32(name="FunctionStartAddress", comment="virtual address of function")
        yield Offset32(name="FunctionInfoOffset", base=self.pclntable_offset, comment="offset to GOPCLnFunctionInfo, relative to start of GOPCLnHeader structure")


class GOPCLnFunctionInfoPointer64(Struct):

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

    def parse(self):
        yield Va64(name="FunctionStartAddress", comment="virtual address of function")
        yield Offset64(name="FunctionInfoOffset", base=self.pclntable_offset, comment="offset to GOPCLnFunctionInfo, relative to start of GOPCLnHeader structure")


class GOPCLnFunctionInfoPointer_118(Struct):

    def __init__(self, text_start, pclntable_offset, **args):
        Struct.__init__(self, **args)
        self.pclntable_offset = pclntable_offset
        self.text_start = text_start

    def parse(self):
        yield Offset32(name="FunctionStartAddress", base=self.text_start, comment="virtual address of function, relative to start of .text section")
        yield Offset32(name="FunctionInfoOffset", base=self.pclntable_offset, comment="offset to GOPCLnFunctionInfo, relative to start of GOPCLnHeader structure")

def bytes2re(s):
    return "".join(map(lambda x: f"\\x{x:02x}", s))


def parse_golang_structures(analyzer, is64=False, pcln_hdr_offset = None, va2off = lambda x: x, off2va = lambda x:x):

    # search pcln header
    force_118 = False
    last_physical_offset = 0
    first_executable_region_va = None
    pclntab_section_offset = None
    for r in analyzer.regions:
        if first_executable_region_va is None and r.x and r.foff:
            first_executable_region_va = r.va
        if r.fsize and r.foff + r.fsize > last_physical_offset:
            last_physical_offset = r.foff + r.fsize
        if r.name.startswith(".gopclntab"):
            pclntab_section_offset = r.foff
        elif r.name == ".text":
            first_executable_region_va = r.va
    pcln_header = None
    if pcln_hdr_offset is None:
        pcln_hdr_offset, _ = analyzer.search(r"[\xF0\xF1\xFA\xFB]\xFF\xFF\xFF\x00\x00.[\x04\x08]", 0, last_physical_offset)
    if pcln_hdr_offset is None and first_executable_region_va is not None and first_executable_region_va > 0:
        # extra aggressive heuristic for obfuscated Golang 118+ programs, e.g. 83654c500c68418142e43b31ebbec040d9d36cfbbe08c7b9b3dc90fabc14801a
        if is64:
            pattern = r"\x08...\x00\x00\x00\x00\x00...\x00\x00\x00\x00\x00" + bytes2re(struct.pack("<Q", first_executable_region_va)) + r"(?:...\x00\x00\x00\x00\x00){4}"
        elif first_executable_region_va <= 0xffffffff:
            pattern = r"\x04...\x00...\x00" + bytes2re(struct.pack("<I", first_executable_region_va)) + r"(?:...\x00){4}"
        else:
            return
        pat_offset, _ = analyzer.search(pattern, 0, last_physical_offset)
        if pat_offset is not None:
            pcln_hdr_offset = pat_offset - 7
            force_118 = True
    if pcln_hdr_offset is None:
        pcln_hdr_offset = pclntab_section_offset
        force_118 = True
    if pcln_hdr_offset is None:
        return

    is120 = analyzer.read(pcln_hdr_offset, 1) == b"\xF1"
    is118 = analyzer.read(pcln_hdr_offset, 1) == b"\xF0" or force_118
    is116 = analyzer.read(pcln_hdr_offset, 1) == b"\xFA"
    is114 = analyzer.read(pcln_hdr_offset, 1) == b"\xFB"
    analyzer.jump(pcln_hdr_offset)
    if is116 or is118 or is120:
        if is64:
            if is116:
                pcln_header = yield GOPCLnHeader64_116(name="GO.PcLnHeader", category=Type.HEADER)
            elif is118 or is120:
                pcln_header = yield GOPCLnHeader64_118(name="GO.PcLnHeader", category=Type.HEADER)
        else:
            if is116:
                pcln_header = yield GOPCLnHeader32_116(name="GO.PcLnHeader", category=Type.HEADER)
            elif is118 or is120:
                pcln_header = yield GOPCLnHeader32_118(name="GO.PcLnHeader", category=Type.HEADER)
        analyzer.jump(pcln_header["FunctionTable"] + pcln_header.offset)
        function_info_base = pcln_header["FunctionTable"] + pcln_header.offset
    elif is114:
        pcln_header = yield GOPCLnHeader(name="GO.PcLnHeader", category=Type.HEADER)
        function_info_base = pcln_header.offset
        # function table directly follows pcln header

    if pcln_header is None:
        return

    # parse functions
    if is118 or is120:
        textstart = pcln_header["TextStart"] 
        if textstart == 0:
            text_start_off = va2off(first_executable_region_va)
        else:
            text_start_off = va2off(textstart)
            if text_start_off is None:
                raise FatalError("Cannot find TextStart offset")
        pclnfuninfoclass = GOPCLnFunctionInfoPointer_118(text_start_off, function_info_base)
    else:
        if pcln_header["PointerSize"] == 8:
            pclnfuninfoclass = GOPCLnFunctionInfoPointer64(function_info_base)
        else:
            pclnfuninfoclass = GOPCLnFunctionInfoPointer32(function_info_base)
    yield Array(pcln_header["NumberOfFunctions"], pclnfuninfoclass, name="GO.PcLnFunctionInfoPointers", category=Type.DEBUG)
    end_of_function_pointers = analyzer.tell()

    # layout
    pcln_header_va = off2va(pcln_header.offset)
    if pcln_header_va is None:
        return
    if is64:
        ref = struct.pack("<Q", pcln_header_va)
    else:
        ref = struct.pack("<I", pcln_header_va)
    refreg = bytes2re(ref)
    p = 0
    while p < analyzer.size():
        symtab_off, sz = analyzer.search(refreg, p)
        if not sz:
            return
        if is64:
            if is116 or is118 or is120:
                nxt, = struct.unpack("<Q", analyzer.read(symtab_off + 8, 8))
            elif is114:
                nxt, = struct.unpack("<Q", analyzer.read(symtab_off + 24, 8))
        else:
            if is116 or is118 or is120:
                nxt, = struct.unpack("<I", analyzer.read(symtab_off + 4, 4))
            elif is114:
                nxt, = struct.unpack("<I", analyzer.read(symtab_off + 12, 4))
        if nxt > pcln_header_va and nxt - pcln_header_va <= 0x60:
            break
        p = symtab_off + 1
    analyzer.jump(symtab_off)
    if is64:
        if is116:
            layout = yield GOLayoutTable64_116(name="GO.LayoutTable", category=Type.HEADER)
        elif is118 or is120:
            layout = yield GOLayoutTable64_118(name="GO.LayoutTable", category=Type.HEADER)
        elif is114:
            layout = yield GOLayoutTable64(name="GO.LayoutTable", category=Type.HEADER)
    else:
        if is116:
            layout = yield GOLayoutTable32_116(name="GO.LayoutTable", category=Type.HEADER)
        elif is118 or is120:
            layout = yield GOLayoutTable32_118(name="GO.LayoutTable", category=Type.HEADER)
        elif is114:
            layout = yield GOLayoutTable32(name="GO.LayoutTable", category=Type.HEADER)

    if "PcLnSize" in layout:
        rest = pcln_header.offset + layout["PcLnSize"] - end_of_function_pointers
        if rest > 0:
            analyzer.jump(end_of_function_pointers)
            yield Bytes(rest, name="GO.PcLnFunctionInfo", category=Type.DEBUG)
