calc_bootimg_size.py

#!/usr/bin/env python3

import argparse
import os
import struct


PAGE_SIZE = 4096


def read_int32(file):
    return struct.unpack('I', file.read(4))[0]


def align_to(offset, page_size):
    padding = (page_size - (offset & (page_size - 1))) & (page_size - 1)
    return offset + padding


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('image', help='Boot image path')

    return parser.parse_args()


def main():
    args = parse_args()

    with open(args.image, 'rb') as f:
        magic = f.read(8)
        if magic != b'ANDROID!':
            raise Exception('Not a boot image')

        kernel_size = read_int32(f)
        ramdisk_size = read_int32(f)

        # Skip OS version and patch level field
        f.seek(4, os.SEEK_CUR)

        header_size = read_int32(f)

        # Skip reserved
        f.seek(16, os.SEEK_CUR)

        version = struct.unpack('I', f.read(4))[0]
        if version not in (3, 4):
            raise Exception(f'Header version {version} not supported')

        # The boot image is header + kernel + ramdisk, where each section has
        # padding at the end to align it to the page size.
        size = header_size
        size = align_to(size, PAGE_SIZE)

        size += kernel_size
        size = align_to(size, PAGE_SIZE)

        size += ramdisk_size
        size = align_to(size, PAGE_SIZE)

        print(size)


if __name__ == '__main__':
    main()