from filetypes.base import *
import malcat


class UImageHeader(Struct):

    def parse(self):
        yield UInt32BE(name="Signature", comment="magic")
        yield UInt32BE(name="HeaderCrc32", comment="header CRC32")
        yield TimestampBE(name="Timestamp", comment="image creation date")
        yield UInt32BE(name="ImageSize", comment="size of image")
        yield Va32BE(name="LoadAddress", comment="virtual address where the image should be loaded")
        yield Va32BE(name="EntryAddress", comment="entrypoint")
        yield UInt32BE(name="DataCrc32", comment="CRC32 of the image data")
        yield UInt8(name="OsType", values=[
            ("IH_OS_INVALID",               0),
            ("IH_OS_OPENBSD",               1),
            ("IH_OS_NETBSD",                2),
            ("IH_OS_FREEBSD",               3),
            ("IH_OS_4_4BSD",                4),
            ("IH_OS_LINUX",                 5),
            ("IH_OS_SVR4",                  6),
            ("IH_OS_ESIX",                  7),
            ("IH_OS_SOLARIS",               8),
            ("IH_OS_IRIX",                  9),
            ("IH_OS_SCO",                   10),
            ("IH_OS_DELL",                  11),
            ("IH_OS_NCR",                   12),
            ("IH_OS_LYNXOS",                13),
            ("IH_OS_VXWORKS",               14),
            ("IH_OS_PSOS",                  15),
            ("IH_OS_QNX",                   16),
            ("IH_OS_U_BOOT",                17),
            ("IH_OS_RTEMS",                 18),
            ("IH_OS_ARTOS",                 19),
            ("IH_OS_UNITY",                 20),
            ("IH_OS_INTEGRITY",             21),
            ("IH_OS_OSE",                   22),
            ("IH_OS_PLAN9",                 23),
            ("IH_OS_OPENRTOS",              24),
            ("IH_OS_ARM_TRUSTED_FIRMWARE",  25),
            ("IH_OS_TEE",                   26),
            ("IH_OS_OPENSBI",               27),
            ("IH_OS_EFI",                   28),
        ])
        yield UInt8(name="Architecture", values=[
            ("IH_ARCH_INVALID",             0),
            ("IH_ARCH_ALPHA",               1),
            ("IH_ARCH_ARM",                 2),
            ("IH_ARCH_I386",                3),
            ("IH_ARCH_IA64",                4),
            ("IH_ARCH_MIPS",                5),
            ("IH_ARCH_MIPS64",              6),
            ("IH_ARCH_PPC",                 7),
            ("IH_ARCH_S390",                8),
            ("IH_ARCH_SH",                  9),
            ("IH_ARCH_SPARC",               10),
            ("IH_ARCH_SPARC64",             11),
            ("IH_ARCH_M68K",                12),
            ("IH_ARCH_NIOS",                13),
            ("IH_ARCH_MICROBLAZE",          14),
            ("IH_ARCH_NIOS2",               15),
            ("IH_ARCH_BLACKFIN",            16),
            ("IH_ARCH_AVR32",               17),
            ("IH_ARCH_ST200",               18),
            ("IH_ARCH_SANDBOX",             19),
            ("IH_ARCH_NDS32",               20),
            ("IH_ARCH_OPENRISC",            21),
            ("IH_ARCH_ARM64",               22),
            ("IH_ARCH_ARC",                 23),
            ("IH_ARCH_X86_64",              24),
            ("IH_ARCH_XTENSA",              25),
            ("IH_ARCH_RISCV",               26),
        ])
        yield UInt8(name="ImageType", values=[
            ("IH_TYPE_INVALID",             0),
            ("IH_TYPE_STANDALONE",          1),
            ("IH_TYPE_KERNEL",              2),
            ("IH_TYPE_RAMDISK",             3),
            ("IH_TYPE_MULTI",               4),
            ("IH_TYPE_FIRMWARE",            5),
            ("IH_TYPE_SCRIPT",              6),
            ("IH_TYPE_FILESYSTEM",          7),
            ("IH_TYPE_FLATDT",              8),
            ("IH_TYPE_KWBIMAGE",            9),
            ("IH_TYPE_IMXIMAGE",            10),
            ("IH_TYPE_UBLIMAGE",            11),
            ("IH_TYPE_OMAPIMAGE",           12),
            ("IH_TYPE_AISIMAGE",            13),
            ("IH_TYPE_KERNEL_NOLOAD",       14),
            ("IH_TYPE_PBLIMAGE",            15),
            ("IH_TYPE_MXSIMAGE",            16),
            ("IH_TYPE_GPIMAGE",             17),
            ("IH_TYPE_ATMELIMAGE",          18),
            ("IH_TYPE_SOCFPGAIMAGE",        19),
            ("IH_TYPE_X86_SETUP",           20),
            ("IH_TYPE_LPC32XXIMAGE",        21),
            ("IH_TYPE_LOADABLE",            22),
            ("IH_TYPE_RKIMAGE",             23),
            ("IH_TYPE_RKSD",                24),
            ("IH_TYPE_RKSPI",               25),
            ("IH_TYPE_ZYNQIMAGE",           26),
            ("IH_TYPE_ZYNQMPIMAGE",         27),
            ("IH_TYPE_ZYNQMPBIF",           28),
            ("IH_TYPE_FPGA",                29),
            ("IH_TYPE_VYBRIDIMAGE",         30),
            ("IH_TYPE_TEE",                 31),
            ("IH_TYPE_FIRMWARE_IVT",        32),
            ("IH_TYPE_PMMC",                33),
            ("IH_TYPE_STM32IMAGE",          34),
            ("IH_TYPE_SOCFPGAIMAGE_V1",     35),
            ("IH_TYPE_MTKIMAGE",            36),
            ("IH_TYPE_IMX8MIMAGE",          37),
            ("IH_TYPE_IMX8IMAGE",           38),
            ("IH_TYPE_COPRO",               39),
            ("IH_TYPE_SUNXI_EGON",          40),
        ])
        yield UInt8(name="CompressionType", values=[
	    ("IH_COMP_NONE",	            0),
            ("IH_COMP_GZIP",                1),
            ("IH_COMP_BZIP2",               2),
            ("IH_COMP_LZMA",                3),
            ("IH_COMP_LZO",                 4),
            ("IH_COMP_LZ4",                 5),
            ("IH_COMP_ZSTD",                6),
        ])
        yield StringUtf8(32, name="Name")
        



class UImageAnalyzer(FileTypeAnalyzer):
    category = malcat.FileType.ARCHIVE
    name = "UImage"
    regexp = r"\x27\x05\x19\x56"

    def unpack(self, vfile, password=None):
        hdr = self["Header"]
        compress = hdr.CompressionType.enum
        data = self.read(hdr.offset + hdr.size, hdr["ImageSize"])
        if compress == "IH_COMP_NONE":
            return data
        elif compress == "IH_COMP_GZIP":
            import zlib
            return zlib.decompress(data, wbits=31)
        elif compress == "IH_COMP_BZIP2":
            import bz2
            return bz2.decompress(data)
        elif compress == "IH_COMP_LZMA":
            import lzma
            return lzma.decompress(data)
        elif compress == "IH_COMP_LZO":
            try:
                import lzo
            except ImportError:
                raise ImportError("You need to install python-lzo library first. On Windows I suggest that you use some pre-compiled wheel")
            # TODO: with or without header ?
            return lzo.decompress(data)
        raise NotImplementedError("Unsupported compression algorithm {}".format(compress))

    def parse(self, hint):
        hdr = yield UImageHeader(name="Header")

        # check header crc32
        import zlib
        hdr_data = bytearray(hdr.bytes)
        # exclude header crc32 field
        hdr_data[hdr.HeaderCrc32.offset:hdr.Timestamp.offset] = b"\x00\x00\x00\x00"
        if zlib.crc32(hdr_data) != hdr["HeaderCrc32"]:
            raise FatalError("Invalid header crc")

        self.confirm()      # from this point, all parsing errors are just warnings

        self.add_metadata("Creation date", hdr["Timestamp"].strftime("%Y-%m-%d %H:%M:%S"))
        if hdr["Name"]:
            self.add_metadata("Image name", hdr["Name"])

        data_size = hdr["ImageSize"]
        if data_size > self.remaining():
            raise FatalError("Truncated image")

        # uimage has a single section
        self.add_section("content", self.tell(), data_size, hdr["LoadAddress"], data_size, r=True, x=True)

        # uimage has a single "member":
        self.add_file("<content>", data_size, "unpack")
