#!/usr/bin/env python3

# -*- coding: utf-8 -*-

# Source: https://github.com/csftech/Xiaomi-OpenWrt-firmware-toolkit
# Author: @csftech (Sheng-Fu Chang)
# MIT Licensed

import argparse
import binascii
import ctypes
import hashlib
import logging
import os
import time


class ImageHeader(ctypes.Structure):
    _fields_ = [
        ("magic", ctypes.c_uint),
        ("signature_offset", ctypes.c_uint),
        ("crc32_checksum", ctypes.c_uint),
        ("file_type", ctypes.c_ushort),
        ("model", ctypes.c_ushort),
        ("segment_offsets", ctypes.c_uint * 8),
    ]


class SegmentHeader(ctypes.Structure):
    _fields_ = [
        ("magic", ctypes.c_uint),
        ("flash_address", ctypes.c_uint),
        ("length", ctypes.c_uint),
        ("partition", ctypes.c_uint),
        ("segment_name", ctypes.c_char * 32),
    ]


class Firmware:
    def __init__(self, path):
        self.path = path
        self.image_header = ImageHeader()
        self.fd = open(self.path, "rb")

    def verify(self, ignore_hash=False):
        logging.info("[Jobs] Verifying firmware image...")
        assert self.fd.readinto(self.image_header) == ctypes.sizeof(self.image_header)

        # magic
        logging.info(f"firmware magic: {hex(self.image_header.magic)}")
        assert self.image_header.magic in [0x31524448, 0x32524448]

        if self.image_header.magic == 0x32524448:
            logging.info("NOTE: firmware seems to be encrypted?")

        # signature
        self.fd.seek(self.image_header.signature_offset)
        signature_length = int.from_bytes(self.fd.read(16), "little")
        signature = self.fd.read()
        logging.info(
            f"firmware signature_offset: {hex(self.image_header.signature_offset)}"
        )
        logging.info(f"firmware signature_length: {signature_length}")
        logging.info(f"firmware signature: {signature.hex()}")
        assert len(signature) == signature_length

        # crc32
        self.fd.seek(12)
        logging.info(
            f"firmware crc32_checksum: {hex(self.image_header.crc32_checksum)}"
        )
        computed_checksum = ~binascii.crc32(self.fd.read()) & 0xFFFFFFFF
        logging.info(f"computed crc32_checksum: {hex(computed_checksum)}")
        assert self.image_header.crc32_checksum == computed_checksum

        # md5
        self.fd.seek(0)
        m = hashlib.md5()
        m.update(self.fd.read())
        hash = m.hexdigest()
        logging.info(f"computed md5 hash: {hash}")
        filename_hash = os.path.basename(self.path).split("_")[-2]
        try:
            assert hash[(len(filename_hash) * -1) :] == filename_hash
        except AssertionError:
            if ignore_hash:
                logging.warning(
                    f'Warning: hash "{filename_hash}" does not match expected!'
                )
                pass
            else:
                raise

        self.fd.seek(0)
        return True

    def extract(self, dest: str = None):
        logging.info("[Jobs] Extracting firmware...")
        current_time = time.strftime("%Y%m%d_%H%M%S", time.localtime())
        self.dest_dir = os.path.join(
            dest, f'{os.path.basename(self.path).replace(".bin", "")}_{current_time}'
        )
        if dest is not None:
            self.dest_dir = dest
        os.mkdir(self.dest_dir)
        logging.info(f"create destination directory: {self.dest_dir}")

        for address in self.image_header.segment_offsets:
            if address:
                self.fd.seek(address)
                segment_header = SegmentHeader()
                assert self.fd.readinto(segment_header) == ctypes.sizeof(segment_header)

                # extract segment
                with open(
                    os.path.join(self.dest_dir, segment_header.segment_name.decode()),
                    "wb",
                ) as s:
                    s.write(self.fd.read(segment_header.length))
                    logging.info(
                        f'extracting segment: {segment_header.segment_name.decode("ascii")}'
                    )


def run(path, extract=False, dest=None, ignore_hash=False):
    logging.info(f"[MSG] Input file: {path}")
    firmware = Firmware(path)
    if firmware.verify(ignore_hash):
        logging.info(
            "[Jobs] Verification success: it's a genuine firmware from Xiaomi."
        )

        if extract:
            firmware.extract(dest)
            logging.info("[Jobs] Extraction complete.")


if __name__ == "__main__":
    logging.basicConfig(format="%(asctime)s - %(message)s", level=logging.INFO)

    parser = argparse.ArgumentParser()
    parser.add_argument(
        "-d",
        "--dest",
        help="Destination directory to store extracted files, default=current working directory",
        action="store",
    )
    parser.add_argument(
        "-i",
        "--ignore-hash",
        help="Do not fail if md5 hash does not match",
        action="store_true",
    )
    parser.add_argument("-e", "--extract", help="Input your firmware", action="store")
    parser.add_argument("-s", "--show", help="Show firmware info", action="store")
    args = parser.parse_args()

    if args.extract:
        firmware_path = os.path.abspath(args.extract)
        dest_path = os.getcwd()
        if args.dest:
            dest_path = args.dest
        run(firmware_path, extract=True, dest=dest_path, ignore_hash=args.ignore_hash)
    elif args.show:
        firmware_path = os.path.abspath(args.show)
        run(firmware_path, ignore_hash=args.ignore_hash)
    else:
        parser.print_help()