linux-exploit-dev-env/build.py

#!/usr/bin/env python3

	# @echo 'Generic Targets:'
	# @echo '   bl                            - builds U-Boot and TF-A and copies them to the bl output directory'
	# @echo '   tfa                           - builds Arm Trusted Firmware-A (BL1, BL2, BL31)'
	# @echo '   bl33/uboot                    - builds U-Boot (BL33)'
	# @echo '   tools-vm                      - builds linux/tools/vm'

# TODO: Shallow clone

import os
import argparse
from argparse import ArgumentParser, _HelpAction
from enum import Enum
import subprocess

class Arch(str, Enum):
    X86_64 = "x86_64"
    ARM64 = "arm64"

def handle_build(args, env):
    subprocess.run(["make", args.target], env=env)


def handle_run(args, env):
    if args.gdb:
        env['GDB'] = '1'
    if args.mem:
        env['MEM'] = args.mem
    if args.cpu:
        env['CPU'] = str(args.cpu)
    if args.extra_qemu_args:
        env["QEMU_EXTRA_ARGS"] = args.extra_qemu_args
    if args.extra_kernel_cmdline:
        env["QEMU_EXTRA_KERNEL_CMDLINE"] = args.extra_kernel_cmdline

    subprocess.run(["make", "run"], env=env)


def handle_clean(args, env):
    if args.all:
        subprocess.run(["make", "clean"], env=env)
    else:
        subprocess.run(["make", f"{args.target}_clean"], env=env)


def handle_rootfs(args, env):
    print(args)

class HelpAction(argparse.Action):
    def __init__(self, option_strings, subparsers, parent_parser, *args, **kwargs):
        self.subparsers = subparsers
        self.parent_parser = parent_parser
        super(HelpAction, self).__init__(option_strings=option_strings, *args, **kwargs)

    def __call__(self, parser, namespace, values, option_string=None):
        parser.print_help()

        f = self.parent_parser._get_formatter()
        f.add_arguments(self.parent_parser._actions)
        lines = f.format_help().split('\n')
        for l in lines:
            print("  " + l)

        for k, v in self.subparsers.choices.items():
            f = v._get_formatter()
            f.start_section(f"{k} options")
            actions = [a for a in v._actions if a not in self.parent_parser._actions and not isinstance(a, _HelpAction)]
            f.add_arguments(actions)
            f.end_section()
            print(f.format_help())

        exit(0)


def parse_args():
    parser = ArgumentParser(description='TODO', add_help=False)
    subparsers = parser.add_subparsers()
    parent_parser = ArgumentParser(add_help=False)

    parser.add_argument('-h', '--help', action=HelpAction, subparsers=subparsers, parent_parser=parent_parser, nargs=0)

    # TODO: How to add parent_parser defaults?
    parent_parser.add_argument('-v', '--verbose', action='store_true')
    parent_parser.add_argument('-a', '--arch', type=Arch, choices=[a.value for a in Arch], default=Arch.X86_64.value)
    parent_parser.add_argument('--ack', action='store_true')
    parent_parser.add_argument('-t', '--target', choices=['linux', 'ack'], default='linux')


    build_parser = subparsers.add_parser('build', parents=[parent_parser], help="Build a kernel image")
    build_parser.set_defaults(func=handle_build)
    # TODO: mutually exclusive
    build_parser.add_argument('-c', '--defconfig', action='store_true')
    build_parser.add_argument('-m', '--modules', action='store_true')

    run_parser = subparsers.add_parser('run', parents=[parent_parser], help="Run a kernel")
    run_parser.set_defaults(func=handle_run)

    run_parser.add_argument('-g', '--gdb', action='store_true')
    run_parser.add_argument('-c', '--cpu', type=int)
    run_parser.add_argument('-m', '--mem')
    run_parser.add_argument('--rootfs')
    run_parser.add_argument('--initramfs')
    run_parser.add_argument('--extra-qemu-args')
    run_parser.add_argument('--extra-kernel-cmdline')

    clean_parser = subparsers.add_parser('clean', parents=[parent_parser], help="Clean built artifacts")
    clean_parser.set_defaults(func=handle_clean)
    clean_parser.add_argument('--all', action='store_true')

    rootfs_parser = subparsers.add_parser('rootfs')
    rootfs_parser.set_defaults(func=handle_rootfs)

    rootfs_subparser = rootfs_parser.add_subparsers(dest="rootfs_command")
    mount_subparser = rootfs_subparser.add_parser('mount')
    unmount_subparser = rootfs_subparser.add_parser('unmount')
    initramfs_subparser = rootfs_subparser.add_parser('initramfs')
    image_subparser = rootfs_subparser.add_parser('image')

    return parser.parse_args()


def main():
    args = parse_args()
    print(args)
    env = os.environ.copy()
    # env['ARCH'] = args.arch
    # if args.verbose:
    #     env['VERBOSE'] = '1'
    # if args.target == "ack":
    #     env['ACK'] = '1'

    args.func(args, env)

if __name__ == '__main__':
    main()