import argparse
import pychrysalide
from pychrysalide.analysis.contents import FileContent
from pychrysalide.format.bootimg import BootImgFormat
import struct


def _value_to_string(val):
    """Convert a field into a string."""

    encodings = [ 'ascii', 'utf-8', 'utf-16' ]

    desc = None

    for enc in encodings:

        try:
            desc = '"%s"' % val.decode(enc)
        except:
            desc = None

        if desc:
            break

    if not(desc):

        count = val.find(b'\x00')

        if count == -1:
            count = len(val)

        desc = '%u byte(s)' % count

    return desc


def _value_to_size(val):
    """Convert a field into a human-readable size."""

    unit = 0

    units = [ '', 'Kb', 'Mb', 'Gb' ]

    for i in range(3):
        if val > 1024:
            val /= 1024
            unit += 1

    if val == int(val):
        return '%u%s' % (val, units[unit])
    else:
        return '%.1f%s' % (val, units[unit])


def _value_to_id(val):
    """Convert a field into an identifier."""

    values = struct.unpack('<IIIIIIII', val)

    return ' '.join(['0x%08x' % v for v in values])


def show_info(fmt):
    """Display all available information about a given Boot image."""

    hdr = fmt.header

    items = [

        [ 'magic', _value_to_string(hdr.magic) ],
        [ 'kernel_size', _value_to_size(hdr.kernel_size) ],
        [ 'kernel_addr', hex(hdr.kernel_addr) ],

        [ 'ramdisk_size', _value_to_size(hdr.ramdisk_size) ],
        [ 'ramdisk_addr', hex(hdr.ramdisk_addr) ],

        [ 'second_size', _value_to_size(hdr.second_size) ],
        [ 'second_addr', hex(hdr.second_addr) ],

        [ 'tags_addr', hex(hdr.tags_addr) ],
        [ 'page_size', _value_to_size(hdr.page_size) ],
        [ 'header_version', hdr.header_version ],
        [ 'os_version', hex(hdr.os_version) ],
        [ 'name', _value_to_string(hdr.name) ],
        [ 'cmdline', _value_to_string(hdr.cmdline) ],
        [ 'id', _value_to_id(hdr.id) ],
        [ 'extra_cmdline', _value_to_string(hdr.extra_cmdline) ],

        [ 'recovery_dtbo_size', _value_to_size(hdr.recovery_dtbo_size) ],
        [ 'recovery_dtbo_offset', hex(hdr.recovery_dtbo_offset) ],
        [ 'header_size', _value_to_size(hdr.header_size) ],

    ]

    length = 0

    for key, _ in items:
        if (len(key) + 1) > length:
            length = len(key) + 1

    line = '  {:<%u} {:}' % length

    title = ' Boot image header'

    print()
    print(title)
    print('=' * len(title))

    for key, val in items:
        print(line.format(key + ':', val))

    print()


if __name__ == '__main__':
    """Script entry point."""

    parser = argparse.ArgumentParser()

    parser.add_argument('-i', '--info',    help='display information about the boot image', action='store_true')
    parser.add_argument('-k', '--kernel',  help='dump the kernel embedded in the boot image', action='store_true')
    parser.add_argument('-r', '--ramdisk', help='dump the ramdisk embedded in the boot image', action='store_true')
    parser.add_argument('filename',        help='path to the Android boot image', metavar='boot.img')

    args = parser.parse_args()

    cnt = FileContent(args.filename)

    if cnt is None:
        sys.exit('No content to load!')

    fmt = BootImgFormat(cnt)

    status = fmt.analyze()

    if not(status):
        sys.exit('Bad boot image!')

    if args.info:
        show_info(fmt)

    if args.kernel:

        if fmt.kernel:

            with open('kernel.bin', 'wb') as fd:
                fd.write(fmt.kernel.data)

        else:

            print('No kernel to extract!')

    if args.ramdisk:

        if fmt.ramdisk:

            with open('ramdisk.bin', 'wb') as fd:
                fd.write(fmt.ramdisk.data)

        else:

            print('No ramdisk to extract!')