#!/usr/bin/env python3

"""egoboost.py: Build U-Boot micro SD card images for ego machines

See also:
- https://blog.pcfe.net/hugo/posts/2019-01-27-fedora-29-on-odroid-hc2/
- https://blog.christophersmart.com/2018/04/28/fedora-on-odroid-hc1-mini-nas-armv7/
"""


import argparse
import os
import shutil
import stat
import subprocess
import sys
import textwrap

from urllib import request


# Constants that we get from the Jinja2 template
DEFAULT_EGOBOOST_LIB_DIRECTORY = "{{ egoboost_lib_directory }}"


def resolvepath(*path):
    if isinstance(path, tuple):
        path = list(path)
    elif not isinstance(path, str):
        path = [path]
    return os.path.realpath(os.path.normpath(os.path.expanduser(os.path.join(*path))))


def is_block_device(path):
    return stat.S_ISBLK(os.stat(path).st_mode)


def dd_dsync(infile, outfile, seek, ddargs=None):
    ddargs = ddargs or []
    return subprocess.run(
        ['dd', 'iflag=dsync', 'oflag=dsync', f'if={infile}', f'of={outfile}', f'seek={seek}'] + ddargs,
        check=True)


def main(*args, **kwargs):
    parser = argparse.ArgumentParser(
        "egoboost - Build U-Boot micro SD card images for ego machines")
    parser.add_argument(
        "--lib-directory", default=DEFAULT_EGOBOOST_LIB_DIRECTORY,
        help="Location of binary blobs from hardkernel")
    parser.add_argument(
        '--u-boot', choices=['hardkernel', 'rpm'], required=True,
        help="Use Das U-Boot from either HardKernel's repo or the 'uboot-images-armv7' RPM")
    parser.add_argument(
        'memcard', help='Device path to the memory card (whole card, NOT partition)')
    parsed = parser.parse_args()

    if not is_block_device(parsed.memcard):
        raise Exception(f"{parsed.memcard} is not a block device, cannot continue")

    partition_one = f'{parsed.memcard}1'

    firmware = {
        'bl1': resolvepath(parsed.lib_directory, 'bl1.bin.hardkernel'),
        'bl2': resolvepath(parsed.lib_directory, 'bl2.bin.hardkernel.720k_uboot'),
        'tz': resolvepath(parsed.lib_directory, 'tzsw.bin.hardkernel'),
        'u_boot_hardkernel': resolvepath(parsed.lib_directory, 'u-boot.bin.hardkernel'),
        'u_boot_rpm': '/usr/share/uboot/odroid-xu3/u-boot.bin',
    }
    firmware['u_boot'] = firmware[f'u_boot_{parsed.u_boot}']

    signed_bl1_position = 1
    bl2_position = 31
    uboot_position = 63
    tzsw_position = 1503
    uboot_env_position = 2015

    # 1GB in sectors
    # 1024k * 1024k/m * 1024m/g * 1/512sectors/k = 2097152 sectors/GB
    #
    # TODO: Discover sector size from /sys/block/sda/queue/hw_sector_size
    # It's theoretically possible that future SD cards will have a different sector size than 512
    #
    # NOTE: our filesystem is VFAT, which cannot be larger than 2GB, maybe
    onegb_sectors = int(1024 * 1024 * 1024 / 512)

    # This partition dump uses partition type 83, aka Linux ext2, but that doesn't matter
    sfdisk_dump = textwrap.dedent(f"""
        label: dos
        label-id: 0x00000000
        device: {parsed.memcard}
        unit: sectors

        {partition_one} : start=        2048, size=     {onegb_sectors}, type=83
    """.strip())

    subprocess.run(['sfdisk', parsed.memcard], input=sfdisk_dump.encode(), check=True)
    subprocess.run(['mkfs.vfat', '-n', 'boot', partition_one], check=True)

    print("Fusing BL1...")
    dd_dsync(firmware['bl1'], parsed.memcard, signed_bl1_position)
    print("Fusing BL2...")
    dd_dsync(firmware['bl2'], parsed.memcard, bl2_position)
    print("Fusing Das U-Boot...")
    dd_dsync(firmware['u_boot'], parsed.memcard, uboot_position)
    print("Fusing TrustZone...")
    dd_dsync(firmware['bl1'], parsed.memcard, tzsw_position)
    print("Erasing Das U-Boot env...")
    dd_dsync('/dev/zero', parsed.memcard, uboot_env_position, ddargs=['bs=512', 'count=32'])

    print("Das U-boot image is fused successfully.")

    print(f"Eject {parsed.memcard} and insert it again.")


if __name__ == '__main__':
    sys.exit(main(*sys.argv))

