from filetypes.base import *
import malcat
import struct
from filetypes.FAT32 import Fat32Analyzer, align

# https://en.wikipedia.org/wiki/Design_of_the_FAT_file_system

class BootSector(Struct):

    def parse(self):
        yield Bytes(3, name="InitialJump")
        yield String(8, name="OemName")
        yield ExtendedBiosParameter()
        yield Bytes(510 - len(self), name="BootloaderCode")
        yield UInt16(name="Signature")



class ExtendedBiosParameter(Struct):

    def parse(self):
        yield UInt16(name="BytesPerSector")
        yield UInt8(name="SectorsPerCluster")
        yield UInt16(name="ReservedSectors")
        yield UInt8(name="NumberOfFAT")
        yield UInt16(name="MaximumNumberOfRootDirectoryEntries")
        yield UInt16(name="TotalLogicalSectors")
        yield UInt8(name="MediaType", values=[
            ("FixedDisk", 0xf8),
        ])
        yield UInt16(name="SectorsPerFat")
        yield UInt16(name="SectorsPerTrack")
        yield UInt16(name="NumberOfHeads")
        yield UInt32(name="HiddenSectors")
        yield UInt32(name="TotalLogicalSectors")
        yield UInt8(name="DriveNumber")
        yield Unused(1)
        yield String(1, name="ExtendedBootSignature")
        yield UInt32(name="VolumeId")
        yield String(11, name="VolumeLabel")
        yield String(8, name="FileSystemType")


class DirectoryEntry(Struct):

    def parse(self):
        yield String(8, name="ShortFileName")
        yield String(3, name="FileExtension")
        yield BitsField(
            Bit(name="ReadOnly"),
            Bit(name="Hidden"),
            Bit(name="System"),
            Bit(name="VolumeLabel"),
            Bit(name="Directory"),
            Bit(name="Archive"),
            Bit(name="Device"),
            name="FileAttributes")
        yield UInt8(name="ExtraAttributes")
        yield UInt8(name="CreationTimeFine")
        yield DosDateTime(name="CreationTime")
        yield DosDate(name="LastAccess")
        yield UInt16(name="Attributes")
        yield DosDateTime(name="ModificationTime")
        yield UInt16(name="FirstCluster")
        yield UInt32(name="FileSize")



class Fat12Analyzer(Fat32Analyzer):
    category = malcat.FileType.FILESYSTEM
    name = "FAT12"
    regexp = r"\xEB.\x90.{8}\x00\x02[\x01\x02\x04\x08\x10\x20\x40\x80]..\x02....\xf8...{14}\x29.{15}FAT12   .{448}\x55\xaa"

    def iter_cluster_chain(self, cluster):
        seen = set()
        while cluster >= 2 and cluster <= 0xFEF and cluster not in seen:
            seen.add(cluster)
            yield cluster
            if cluster >= len(self.fat):
                raise FatalError("No FAT entry for cluster {:x}".format(cluster))
            cluster = self.fat[cluster]

    def parse_dir(self, cluster, parents=[], size = None, seen=None):
        if seen is None:
            seen = set()
        elif cluster in seen:
            return
        seen.add(cluster)
        for offset, sz in self.iter_file_intervals(cluster, size):
            yield from self.parse_dir_range(offset, sz, parents, seen=seen)
            
    def parse_dir_range(self, offset, sz, parents=[], seen=None):
        self.jump(offset)
        entries = yield Array(sz // 32, DirectoryEntry(), name="Directory")
        long_name_bytes = b""
        for e in entries:
            is_dir = e["FileAttributes"]["Directory"]
            filesize = e["FileSize"]
            if filesize == 0 and not is_dir:
                continue
            if e["FileAttributes"]["VolumeLabel"] and e["FileAttributes"]["System"] and e["FileAttributes"]["Hidden"] and e["FileAttributes"]["ReadOnly"]:
                cur_bytes = self.read(e.offset + 1, 10) + self.read(e.offset + 14, 12) + self.read(e.offset + 28, 4)
                long_name_bytes = cur_bytes + long_name_bytes
            else:
                name = e["ShortFileName"].strip()
                ext = e["FileExtension"].strip()
                if ext:
                    name += "." + ext
                if name and ord(name[0]) == 0xE5:
                    continue
                if long_name_bytes:
                    name = long_name_bytes.decode("utf-16-le").strip()
                    i = name.find("\x00")
                    if i >= 0:
                        name = name[:i]
                    long_name_bytes = b""
                cluster = e["FirstCluster"]
                if is_dir:
                    if cluster > 0:
                        if filesize == 0:
                            filesize = None
                        yield from self.parse_dir(cluster, parents + [name], filesize, seen=seen)
                else:
                    fname = "/".join(parents + [name])
                    self.filesystem[fname] = (cluster, filesize)
                    self.add_file(fname, filesize, "open_file")

    
    def parse(self, hint):
        boot_sector = yield BootSector()
        ebp = boot_sector["ExtendedBiosParameter"]
        self.sector_size = ebp["BytesPerSector"]
        self.cluster_size = ebp["SectorsPerCluster"] * self.sector_size
        self.reserved_space = ebp["ReservedSectors"] * self.sector_size
        self.add_section("Boot", 0, self.reserved_space, x=True)
        eof = boot_sector["ExtendedBiosParameter"]["TotalLogicalSectors"] * self.sector_size
        self.set_eof(eof)

        # read FATs
        self.jump(self.reserved_space)
        fat_size = ebp["SectorsPerFat"] * self.sector_size
        self.add_section("FAT", self.reserved_space, ebp["NumberOfFAT"] * fat_size, discardable=True)
        for i in range(ebp["NumberOfFAT"]):
            fat = yield Bytes(fat_size, name="FAT#{}".format(i), category=Type.FIXUP)
            if not self.fat:
                for i in range(0, len(fat) - 2, 3):
                    first, second = struct.unpack_from("<HB", fat, i)
                    second = (second << 4) | (first >> 12)
                    first = first & 0xfff
                    if first * self.cluster_size > eof or second * self.cluster_size > eof:
                        raise FatalError("Cluster outside FS boundaries")
                    self.fat.append(first)
                    self.fat.append(second)
        self.confirm()

        # read root dir
        if not ebp["MaximumNumberOfRootDirectoryEntries"]:
            raise FatalError("Not a valid root dir")
        num_entries = ebp["MaximumNumberOfRootDirectoryEntries"]
        self.user_space = self.tell() + num_entries * 32
        self.add_section("RootDir", self.tell(), num_entries * 32)
        self.root = yield from self.parse_dir_range(self.tell(), num_entries*32)

        self.add_section("Clusters", self.user_space, eof - self.user_space, r=True, w=True)
