#!/usr/bin/env python3
import json
from pathlib import Path

import click

from trezorlib import firmware, toif


def minimum_header_len(spec):
    spec = spec.copy()
    spec["header_len"] = 512000
    reparsed = firmware.VendorHeader.SUBCON.parse(
        firmware.VendorHeader.SUBCON.build(spec)
    )
    data_length = reparsed._end_offset - reparsed._start_offset
    # data length + 65 for signatures, rounded up to nearest multiple of 512
    return (data_length + 65 + 511) // 512 * 512


@click.command()
@click.argument("specfile", type=click.File("r"))
@click.argument("image", type=click.File("rb"))
@click.argument("outfile", type=click.Path(dir_okay=False, writable=True, path_type=Path))
@click.option("-c", "--check", is_flag=True, help="Check but do not write header.")
@click.option("-q", "--quiet", is_flag=True, help="Do not print anything.")
def build_vendorheader(specfile, image, outfile, check: bool, quiet: bool):
    if quiet:
        echo = lambda *args, **kwargs: None
    else:
        echo = click.echo

    spec = json.load(specfile)
    spec["pubkeys"] = [bytes.fromhex(k) for k in spec["pubkeys"]]
    spec["image"] = toif.ToifStruct.parse(image.read())
    spec["sigmask"] = 0
    spec["signature"] = b"\x00" * 64
    if spec["hw_model"] is None:
        spec["hw_model"] = b"\x00\x00\x00\x00"
    else:
        spec["hw_model"] = spec["hw_model"].encode("ascii")

    min_length = minimum_header_len(spec)
    if "header_len" not in spec:
        spec["header_len"] = min_length
    elif spec["header_len"] < min_length:
        raise click.ClickException(
            f"Specified header_len {spec['header_len']} too low. "
            f"Minimum allowable value is {min_length}."
        )
    elif spec["header_len"] == min_length:
        echo(f"{specfile.name}: Header has correct length.")
    else:
        echo(
            f"{specfile.name}: Extending header ({min_length} bytes) to {spec['header_len']} bytes."
        )

    if spec["header_len"] % 512 != 0:
        raise click.ClickException("Invalid header_len: must be a multiple of 512")

    vh_bytes = firmware.VendorHeader.SUBCON.build(spec)
    if check:
        if not outfile.exists():
            raise click.ClickException(f"Header file {outfile.name} does not exist.")
        outfile_bytes = outfile.read_bytes()
        if outfile_bytes != vh_bytes:
            raise click.ClickException(
                f"Header file {outfile.name} differs from expected header."
            )
    else:
        outfile.write_bytes(firmware.VendorHeader.SUBCON.build(spec))


if __name__ == "__main__":
    build_vendorheader()