#!/usr/bin/env python3
"""Checks an Update File (.ota) the way the device does, on a PC: header, signature, image hash.

Usage: scripts/ota_verify.py <file.ota> [public key, default keys/ota-public.pem]
Exits 0 and prints the version if a device would accept the file.
"""
import hashlib
import os
import struct
import subprocess
import sys
import tempfile

HEADER_SIZE = 160
SIGNED_BYTES = 80


def main():
    if len(sys.argv) < 2:
        sys.exit(__doc__)
    root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    key = sys.argv[2] if len(sys.argv) > 2 else os.path.join(root, "keys", "ota-public.pem")
    data = open(sys.argv[1], "rb").read()
    if len(data) < HEADER_SIZE or data[:8] != b"RORO-OTA":
        sys.exit("not an update file")
    fmt, header_size, image_size = struct.unpack("<HHI", data[8:16])
    if fmt != 1 or header_size != HEADER_SIZE:
        sys.exit("unsupported update format")
    image = data[HEADER_SIZE:]
    if len(image) != image_size:
        sys.exit(f"the header announces {image_size} bytes of image, the file has {len(image)}")
    if hashlib.sha256(image).digest() != data[16:48]:
        sys.exit("image corrupted (hash mismatch)")
    version = data[48:80].split(b"\0")[0].decode()
    (sig_len,) = struct.unpack("<H", data[80:82])
    with tempfile.NamedTemporaryFile() as signed, tempfile.NamedTemporaryFile() as sig:
        signed.write(data[:SIGNED_BYTES])
        signed.flush()
        sig.write(data[82:82 + sig_len])
        sig.flush()
        ok = subprocess.run(["openssl", "dgst", "-sha256", "-verify", key, "-signature", sig.name, signed.name],
                            capture_output=True).returncode == 0
    if not ok:
        sys.exit("bad signature (wrong key)")
    print(f"{sys.argv[1]}: {version}, {image_size} bytes, signature good")


if __name__ == "__main__":
    main()