#!/usr/bin/env python3
# SPDX-License-Identifier: GPL-2.0-or-later

"""Validate and display the bounded structure of an Apple Image3 file."""

import argparse
import json
from pathlib import Path
import struct
import sys


IMG3_HEADER = struct.Struct("<4s4I")
TAG_HEADER = struct.Struct("<4s2I")


class Img3Error(ValueError):
    pass


def fourcc(raw):
    return raw[::-1].decode("ascii", errors="replace")


def scalar(tag, payload):
    if tag == "TYPE" and len(payload) >= 4:
        return fourcc(payload[:4])
    if tag == "VERS" and len(payload) > 4:
        return payload[4:].split(b"\0", 1)[0].decode("ascii", errors="replace")
    if tag in {"CHIP", "BORD", "SEPO", "PROD", "SDOM"} and len(payload) == 4:
        return f"0x{struct.unpack('<I', payload)[0]:08x}"
    if tag == "ECID" and len(payload) == 8:
        return f"0x{struct.unpack('<Q', payload)[0]:016x}"
    if tag == "KBAG" and len(payload) >= 8:
        key_modifier, key_bits = struct.unpack_from("<II", payload)
        expected = 8 + 16 + key_bits // 8
        suffix = "" if len(payload) == expected else f" payload={len(payload)}"
        return (
            f"key_modifier={key_modifier} key_bits={key_bits}" + suffix
        )
    return None


def parse_img3(data):
    if len(data) < IMG3_HEADER.size:
        raise Img3Error("file is shorter than the 20-byte Image3 header")
    magic, full_size, data_size, shsh_offset, image_type = (
        IMG3_HEADER.unpack_from(data)
    )
    if fourcc(magic) != "Img3":
        raise Img3Error(f"bad magic {magic.hex()} (expected Image3)")
    if full_size < IMG3_HEADER.size or full_size > len(data):
        raise Img3Error(
            f"full_size {full_size} is outside the {len(data)}-byte input"
        )
    if data_size > full_size - IMG3_HEADER.size:
        raise Img3Error("data_size extends past full_size")
    if shsh_offset > data_size:
        raise Img3Error("shsh_offset extends past data_size")

    tags = []
    offset = IMG3_HEADER.size
    end = IMG3_HEADER.size + data_size
    while offset < end:
        if end - offset < TAG_HEADER.size:
            raise Img3Error(f"truncated tag header at 0x{offset:x}")
        raw_tag, total_size, payload_size = TAG_HEADER.unpack_from(data, offset)
        tag = fourcc(raw_tag)
        if total_size < TAG_HEADER.size:
            raise Img3Error(f"{tag} at 0x{offset:x} has an invalid total size")
        if payload_size > total_size - TAG_HEADER.size:
            raise Img3Error(
                f"{tag} at 0x{offset:x} has an invalid payload size"
            )
        if total_size > end - offset:
            raise Img3Error(f"{tag} at 0x{offset:x} extends past data_size")
        payload_start = offset + TAG_HEADER.size
        payload = data[payload_start:payload_start + payload_size]
        tags.append({
            "tag": tag,
            "offset": offset,
            "payload_offset": payload_start,
            "total_size": total_size,
            "payload_size": payload_size,
            "padding_size": total_size - TAG_HEADER.size - payload_size,
            "value": scalar(tag, payload),
        })
        offset += total_size
    if offset != end:
        raise Img3Error("tag sizes do not exactly cover data_size")

    return {
        "magic": "Img3",
        "file_size": len(data),
        "full_size": full_size,
        "data_size": data_size,
        "shsh_offset": shsh_offset,
        "image_type": fourcc(struct.pack("<I", image_type)),
        "trailing_size": len(data) - full_size,
        "tags": tags,
    }


def print_human(path, image):
    print(f"{path}: {image['magic']} type={image['image_type']}")
    print(
        f"  file={image['file_size']} full={image['full_size']} "
        f"data={image['data_size']} shsh_offset=0x{image['shsh_offset']:x} "
        f"trailing={image['trailing_size']}"
    )
    for tag in image["tags"]:
        value = f" value={tag['value']}" if tag["value"] is not None else ""
        print(
            f"  0x{tag['offset']:08x} {tag['tag']:<4} "
            f"total={tag['total_size']:<8} data={tag['payload_size']:<8} "
            f"padding={tag['padding_size']}{value}"
        )


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("image", type=Path)
    parser.add_argument("--json", action="store_true", dest="as_json")
    parser.add_argument("--extract-data", type=Path, metavar="PATH")
    args = parser.parse_args()

    data = args.image.read_bytes()
    image = parse_img3(data)
    if args.extract_data:
        data_tags = [tag for tag in image["tags"] if tag["tag"] == "DATA"]
        if len(data_tags) != 1:
            raise Img3Error(f"expected one DATA tag, found {len(data_tags)}")
        tag = data_tags[0]
        start = tag["payload_offset"]
        args.extract_data.write_bytes(data[start:start + tag["payload_size"]])
        print(f"Extracted {tag['payload_size']} bytes to {args.extract_data}")
    if args.as_json:
        print(json.dumps(image, indent=2))
    else:
        print_human(args.image, image)


if __name__ == "__main__":
    try:
        main()
    except (OSError, Img3Error) as error:
        print(f"error: {error}", file=sys.stderr)
        sys.exit(2)