#!/usr/bin/env python3
"""Wraps a firmware image into a signed Update File (.ota). See lib/ota/src/update_parser.h.

Usage: scripts/make_ota.py <firmware.bin> <version> <out.ota> [private key]
Signs with openssl (ECDSA P-256 over SHA-256 of the header's first 80 bytes).
"""
import hashlib
import os
import struct
import subprocess
import sys
import tempfile

HEADER_SIZE = 160
SIGNED_BYTES = 80
MAX_SIGNATURE = 72


def main():
    if len(sys.argv) < 4:
        sys.exit(__doc__)
    image_path, version, out_path = sys.argv[1:4]
    key = sys.argv[4] if len(sys.argv) > 4 else os.environ.get(
        "RORO_OTA_KEY", os.path.expanduser("~/.config/roro9stack/ota-key.pem"))
    if not os.path.exists(key):
        sys.exit(f"No signing key at {key}: run scripts/ota_keygen.sh first.")

    image = open(image_path, "rb").read()
    version_bytes = version.encode()[:31]
    signed = (b"RORO-OTA" + struct.pack("<HHI", 1, HEADER_SIZE, len(image)) +
              hashlib.sha256(image).digest() + version_bytes.ljust(32, b"\0"))
    assert len(signed) == SIGNED_BYTES

    with tempfile.NamedTemporaryFile() as f:
        f.write(signed)
        f.flush()
        signature = subprocess.run(["openssl", "dgst", "-sha256", "-sign", key, f.name],
                                   check=True, capture_output=True).stdout
    if len(signature) > MAX_SIGNATURE:
        sys.exit("unexpected signature size")

    header = signed + struct.pack("<H", len(signature)) + signature
    header = header.ljust(HEADER_SIZE, b"\0")
    with open(out_path, "wb") as out:
        out.write(header + image)
    print(f"{out_path}: {version}, {len(image)} bytes, signed")


if __name__ == "__main__":
    main()