#!/usr/bin/python3
# coding: utf-8
# Copyright (c) 2025 Huawei Technologies Co., Ltd.
# openUBMC is licensed under Mulan PSL v2.
# You can use this software according to the terms and conditions of the Mulan PSL v2.
# You may obtain a copy of Mulan PSL v2 at:
#         http://license.coscl.org.cn/MulanPSL2
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
# See the Mulan PSL v2 for more details.

import os
import sys
import io
from bmcgo.utils.tools import Tools

"""
功 能：打包uboot签名固件
版权信息：华为技术有限公司，版本所有(C) 2025-2025

打包顺序：l0l1 emmc_sub_header rootca.der cms cms.crl rootca.der cms cms.crl uboot.bin
"""

tools = Tools()


class Secion:
    def __init__(self, handler: io.BufferedReader):
        self.handler = handler
        self.offset = 0
        self.length = 0
        self.id = self._read_uint32()
        if self.id == 0x33cc33cc:
            return
        self.offset = self._read_uint32()
        self.length = self._read_uint32()

    def _read_uint32(self):
        # 从流中读取4个字节
        read_len = 4
        ctx = self.handler.read(read_len)
        val = 0
        # 按小端模式生成val
        for i in range(0, read_len):
            # 每个字节向左移(i * 8)位
            val |= (ctx[i] << (i * 8))
        return val


def make_img(uboot_offset, secsions: list[Secion], fw_header: int):
    img = "Hi1711_boot_4096.bin"
    tools.log.info(f"Start make {img}")
    cmd = f"dd if=/dev/zero of={img} bs=1M count=2"
    tools.run_command(cmd)
    cmd = f"dd if=l0l1.bin of={img} conv=notrunc"
    tools.run_command(cmd)
    cmd = f"dd if=u-boot_cms.bin of={img} seek={uboot_offset * 1024} bs=1 count={fw_header} conv=notrunc"
    tools.run_command(cmd)
    # 从emmc_sub_header中提取86个字符存储在输出文件的uboot_offset + 76字节位置
    cmd = f"dd if=emmc_sub_header of={img} seek={uboot_offset * 1024 + 76} bs=1 skip=76 count=84 conv=notrunc"
    tools.run_command(cmd)
    for sec in secsions:
        src = ""
        if sec.id in [0x10, 0x18]:
            src = "rootca.der"
        elif sec.id in [0x11, 0x19]:
            src = "uboot.bin.cms"
        elif sec.id in [0x12, 0x1a]:
            src = "cms.crl"
        elif sec.id == 0x04:
            cmd = f"dd if=uboot.bin of={img} seek={int(uboot_offset + sec.offset / 0x400)} bs=1k conv=notrunc"
            tools.run_command(cmd)
            continue
        else:
            raise Exception(f"Unsupported section {sec.id} get")
        offset = uboot_offset + int(sec.offset / 0x400)
        cmd = f"dd if={src} of={img} seek={offset} bs=1K conv=notrunc"
        tools.run_command(cmd)


if __name__ == "__main__":
    if len(sys.argv) != 2:
        raise Exception(f"Parameter error, usage: {sys.argv[0]} uboot_offset, for example {sys.argv[0]} 768")
    uboot_offset = int(sys.argv[1])

    tools.log.info("Start load emmc_sub_header")
    file_handler = open("emmc_sub_header", "rb")
    # 定位到0x4c位置，该位置存储了section信息
    file_handler.seek(0x4c, os.SEEK_SET)
    sections: list[Secion] = []
    while True:
        sec = Secion(file_handler)
        if sec.id == 0x33cc33cc:
            break
        sections.append(sec)
    file_handler.close()

    file_handler = open("u-boot_cms.bin", "rb")
    # 定位到0x4c位置，该位置存储了section信息
    file_handler.seek(0x4c, os.SEEK_SET)
    fw_sec = None
    while True:
        sec = Secion(file_handler)
        if sec.id == 0x33cc33cc:
            break
        if sec.id == 0x04:
            fw_sec = sec

    header_len = file_handler.tell()
    file_handler.close()

    if not fw_sec:
        raise Exception(f"Find section 0x04 failed")

    make_img(uboot_offset, sections, header_len)
