from filetypes.base import *
import malcat
import struct


# https://www.gnu.org/software/tar/manual/html_node/Standard.html

TAR_BLOCK_SIZE = 512

def align(val, what, down=False):
    if val % what:
        if down:
            val -= val % what
        else:
            val += what - (val % what)
    return val


class HeaderBlock(Struct):

    def parse(self):
        yield String(100, zero_terminated=True, name="Name")
        yield String(8, zero_terminated=True, name="Mode")
        yield String(8, zero_terminated=True, name="UserId")
        yield String(8, zero_terminated=True, name="GroupId")
        sz = yield String(12, zero_terminated=True, name="Size")
        yield String(12, zero_terminated=True, name="ModifiedTime")
        yield String(8, zero_terminated=True, name="Checksum")
        yield UInt8(name="TypeFlag", values = [
            ("AREGTYPE", 0),
            ("REGTYPE", ord("0")),
            ("LNKTYPE", ord("1")),
            ("SYMTYPE", ord("2")),
            ("CHRTYPE", ord("3")),
            ("BLKTYPE", ord("4")),
            ("DIRTYPE", ord("5")),
            ("FIFOTYPE", ord("6")),
            ("CONTTYPE", ord("7")),
        ])
        yield String(100, zero_terminated=True, name="LinkName")
        yield String(6, zero_terminated=True, name="Magic")
        yield String(2, zero_terminated=False, name="Version")
        yield String(32, zero_terminated=True, name="Uname")
        yield String(32, zero_terminated=True, name="Gname")
        yield String(8, zero_terminated=True, name="DevMajor")
        yield String(8, zero_terminated=True, name="DevMinor")
        yield String(155, zero_terminated=True, name="Prefix")
        if len(self) < TAR_BLOCK_SIZE:
            yield Unused(TAR_BLOCK_SIZE - len(self))


class TarAnalyzer(FileTypeAnalyzer):
    category = malcat.FileType.ARCHIVE
    name = "TAR"
    regexp = r"ustar(?:\x0000|  \x00)"

    @classmethod
    def locate(cls, curfile, offset_magic, parent_parser):
        if parent_parser is not None and parent_parser.name == "TAR":
            # no TAR in TAR 
            return None
        if offset_magic >= 257:
            return offset_magic - 257, ""
        return None

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

    def open(self, vfile, password=None):
        block = self.filesystem[vfile.path]
        file_size = int(block["Size"], base=8)
        return self.read(block.offset+len(block), file_size)


    def parse(self, hint):
        num_files = 0
        while self.remaining():
            block = yield HeaderBlock()
            if block["Size"][0] == "\x00" or block["Magic"] != "ustar":
                self.add_section("EOF", block.offset, len(block))
                break
            try:
                file_size = int(block["Size"], base=8)
            except ValueError:
                break
            block_size = align(file_size, TAR_BLOCK_SIZE)
            if block_size:
                name = block["Name"].rstrip("\x00")
                num_files += 1
                self.add_section(name, block.offset, len(block) + block_size)
                self.jump(self.tell() + block_size)
                if block["TypeFlag"] in (0, ord("0")):
                    self.filesystem[name] = block
                    self.add_file(name, file_size, "open")
            elif block["TypeFlag"] != 0x35:
                break
        if not num_files:
            raise FatalError("Empty TAR")
