from filetypes.base import *
import malcat
import json
import os


# https://knifecoat.com/Posts/ASAR+Format+Spec

class AsarHeader(Struct):

    def parse(self):
        yield UInt32(name="DataSize")
        yield UInt32(name="HeaderSize")
        yield UInt32(name="HeaderObjectSize")
        yield UInt32(name="HeaderStringSize")

class AsarAnalyzer(FileTypeAnalyzer):
    category = malcat.FileType.ARCHIVE
    name = "ASAR"
    regexp = r'{"files":{".{0,1024}"size":\d+,'

    @classmethod
    def locate(cls, curfile, offset_magic, parent_parser):
        start_search = max(0, offset_magic - 1024)
        prefix = curfile[start_search:offset_magic]
        x = prefix.rfind(b"\x04\x00\x00\x00")
        if x >= 0:
            return start_search + x, ""
        return None

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

    def open(self, vfile, password=None):
        offset, size = self.filesystem[vfile.path]
        return self.read(offset, size)


    def walk(self, root, path=[]):
        if type(root) != dict:
            return
        for fname, infos in root.get("files", {}).items():
            if type(infos) != dict:
                continue
            if len(infos) == 1 and "files" in infos:
                yield from self.walk(infos, path + [fname])
            elif "offset" in infos:
                name = "/".join(path + [fname])
                offset = int(infos.get("offset", 0))
                size = int(infos.get("size", 0))
                yield name, offset, size

    def parse(self, hint):
        hdr = yield AsarHeader()
        if hdr["HeaderSize"] - hdr["HeaderObjectSize"] != 4:
            raise FatalError("Invalid header")
        self.add_section("listing", self.tell(), hdr["HeaderStringSize"])
        listing_content = yield String(hdr["HeaderSize"] - 8, zero_terminated=True, name="Listing", category=Type.DATA)
        base_offset = self.tell()
        try:
            listing = json.loads(listing_content.rstrip("\x00"))
            if type(listing) != dict:
                raise ValueError("Not a valid listing object")
        except Exception as e:
            raise FatalError(f"Cannot load listing: {e}")
        self.confirm()
        
        for path, offset, size in self.walk(listing):
            if size:
                self.filesystem[path] = (base_offset + offset, size)
                self.add_section(path.split("/")[-1], base_offset + offset, size)
                self.add_file(path, size, "open")
        

                    
            

