from filetypes.base import *
import malcat
import datetime
import struct
import re
from filetypes.MACHO import CPU_MAP


class UniversalBinaryHeader(Struct):

    def parse(self):
        yield UInt32BE(name="Magic")
        yield UInt32BE(name="NumberOfArchitectures")
        

class ArchHeader32(Struct):

    def parse(self):
        cpu = yield UInt32BE(name="CpuType", values=[(v[0], k) for k,v in CPU_MAP.items()])
        stenum = CPU_MAP.get(cpu, ("",[],None))[1]
        st = yield UInt32BE(name="CpuSubType", values=stenum)
        yield Offset32BE(name="Offset")
        yield UInt32BE(name="Size")
        yield UInt32BE(name="Align")
        
class ArchHeader64(Struct):

    def parse(self):
        cpu = yield UInt32BE(name="CpuType", values=[(v[0], k) for k,v in CPU_MAP.items()])
        stenum = CPU_MAP.get(cpu, ("",[],None))[1]
        st = yield UInt32BE(name="CpuSubType", values=stenum)
        yield Offset64BE(name="Offset")
        yield UInt64BE(name="Size")
        yield UInt32BE(name="Align")
        yield Unused(4)




class UniversalBinaryAnalyzer(FileTypeAnalyzer):
    category = malcat.FileType.ARCHIVE
    name = "UMACHO"
    regexp = r"\xca\xfe\xba[\xbe\xbf]\x00\x00\x00[\x01-\x10](?:{})".format("|".join([r"\x{:02x}\x{:02x}\x{:02x}\x{:02x}".format(*struct.pack(">I", x)) for x in CPU_MAP.keys()]))


    def __init__(self):
        FileTypeAnalyzer.__init__(self)
        self.filesystem = {}

    def open(self, vfile, password=""):
        a = self.filesystem.get(vfile.path)
        if a is None:
            raise KeyError(f"Unknown file {vfile.path}")
        return self.read(a["Offset"], a["Size"])

    def parse(self, hint=""):
        hdr = yield UniversalBinaryHeader()
        num = hdr["NumberOfArchitectures"]
        if num == 0 or num > 16:
            raise FatalError("Invalid number of architectures") 
        if hdr["Magic"] == 0xcafebabe:
            archs = yield VariableArray(num, ArchHeader32, name="Architectures")
        else:
            archs = yield VariableArray(num, ArchHeader64, name="Architectures")

        for i, a in enumerate(archs):
            name = a.CpuType.enum
            if not name or name in self.filesystem:
                name = f"Arch[{i}]"
            self.add_section(name, a["Offset"], a["Size"])
            self.add_file(name, a["Size"], "open")
            self.filesystem[name] = a

