add_uncompressed_img_header.py

#!/usr/bin/python3

# SPDX-License-Identifier: GPL-3.0
# Author: Alexey Minnekhanov <alexeymin@postmarketos.org>

# Uncompressed aarch64 kernel image does not include file size
# anywhere inside the file itself; so we have to fake that.
# Append the size as a 32-bit little-endian number as gzip does.

# This works as follows:
# First, "UNCOMPRESSED_IMG" is written to file
# Then, little-endian size of kernel is written
#       in format: bytes are swapped to LE form
# Then kernel image itsef is appended to the result.

# The resulting file can be optionally specified with -o argument.
# If -o is omitted, output will be created in the same folder as the
# source file and will have -hdr suffix appended to its name.

import argparse
import os
import struct
import sys

HEADER='UNCOMPRESSED_IMG'

def main():
    parser = argparse.ArgumentParser(description=f"Append {HEADER} header to uncompressed kernel image")
    parser.add_argument('kernel_image',
                        metavar='KERNEL_IMAGE',
                        type=str,
                        help='path to input kernel Image file')
    parser.add_argument('-o', '--output',
                        metavar='OUTPUT',
                        type=argparse.FileType('wb'),
                        help='path to optional output file with prepended header')
    args = parser.parse_args()

    args.input = open(args.kernel_image, 'rb')
    
    if args.output is None:
        output_fn = f"{args.kernel_image}-hdr"
        args.output = open(output_fn, 'wb')
    else:
        output_fn = args.output.name

    print(f"Creating {HEADER} header for: {args.kernel_image}")
    print(f"    Output file: {output_fn}")

    # determine input file size
    args.input.seek(0, os.SEEK_END)
    insize = args.input.tell()

    print('    Input file size = {} (Hex: {})'.format(insize, hex(insize)))
    # Convert size to little-endian:
    #   < - little endian conversion
    #   I - unsigned int format
    size_le = struct.pack('<I', insize)
    
    # Create output file:
    # 1) write header and kernel size
    print(f"    Writing {HEADER} + {size_le.hex()}...")
    args.output.write(HEADER.encode('utf-8'))
    args.output.write(size_le)
    
    # 2) copy kernel image right after that
    print("    Appending kernel image...")
    args.input.seek(0, os.SEEK_SET)
    ncopied_total = 0
    chunk_of_bytes = args.input.read(4096)
    while len(chunk_of_bytes) != 0:
        args.output.write(chunk_of_bytes)
        # print(len(chunk_of_bytes), ' ', end='')
        ncopied_total += len(chunk_of_bytes)
        chunk_of_bytes = args.input.read(4096)
    
    print(f"    Copied {ncopied_total} bytes.")
    
    if ncopied_total != insize:
        print("    Error - not enough bytes copied!", file=sys.stderr)
    else:
        print('    Success')
    
    args.input.close()
    args.output.flush()
    args.output.close()

if __name__ == "__main__":
    main()