#!/usr/bin/env python3

"""
Whole Flash Image Tag

{
	u32 crc32;
	u32 version;
	u32 chipID;
	u32 flashType;
	u32 flags;
}

CRC32: Ethernet (Poly 0x04C11DB7)

Version:
	0x00005700: Any version
	0x00005731: NAND 1MB data partition
	0x00005732: Normal version

Chip ID:
	Broadcom Chip ID
	0x00006328: BCM6328
	0x00006362: BCM6362
	0x00006368: BCM6368
	0x00063268: BCM63268

Flash Type:
	1: NOR
	2: NAND 16k blocks
	3: NAND 128k blocks
	4: NAND 256k blocks
	5: NAND 512k blocks
	6: NAND 1MB blocks
	7: NAND 2MB blocks

Flags:
	0x00000001: PMC
	0x00000002: Secure BootROM

"""

import argparse
import os
import struct
import binascii


def auto_int(x):
    return int(x, 0)


def create_tag(args, in_bytes):
    # JAM CRC32 is bitwise not and unsigned
    crc = ~binascii.crc32(in_bytes) & 0xFFFFFFFF
    tag = struct.pack(
        ">IIIII",
        crc,
        args.tag_version,
        args.chip_id,
        args.flash_type,
        args.flags,
    )
    return tag


def create_output(args):
    in_st = os.stat(args.input_file)
    in_size = in_st.st_size

    in_f = open(args.input_file, "r+b")
    in_bytes = in_f.read(in_size)
    in_f.close()

    tag = create_tag(args, in_bytes)

    out_f = open(args.output_file, "w+b")
    out_f.write(in_bytes)
    out_f.write(tag)
    out_f.close()


def main():
    global args

    parser = argparse.ArgumentParser(description="")

    parser.add_argument(
        "--input-file",
        dest="input_file",
        action="store",
        type=str,
        help="Input file",
    )

    parser.add_argument(
        "--output-file",
        dest="output_file",
        action="store",
        type=str,
        help="Output file",
    )

    parser.add_argument(
        "--version",
        dest="tag_version",
        action="store",
        type=auto_int,
        help="WFI Tag Version",
    )

    parser.add_argument(
        "--chip-id",
        dest="chip_id",
        action="store",
        type=auto_int,
        help="WFI Chip ID",
    )

    parser.add_argument(
        "--flash-type",
        dest="flash_type",
        action="store",
        type=auto_int,
        help="WFI Flash Type",
    )

    parser.add_argument(
        "--flags", dest="flags", action="store", type=auto_int, help="WFI Flags"
    )

    args = parser.parse_args()

    if not args.flags:
        args.flags = 0

    if (
        (not args.input_file)
        or (not args.output_file)
        or (not args.tag_version)
        or (not args.chip_id)
        or (not args.flash_type)
    ):
        parser.print_help()
    else:
        create_output(args)


main()