ucf

#!/usr/bin/env python3
"""
The main file of unicorefuzz.
This parses the config and provides all the commandline functionality.
"""

from unicornafl import monkeypatch

# Just make sure nothing ever loads an unpatched unicorn in our domain :)
monkeypatch()

import argparse
import os
from typing import Any, Callable, Iterable

from unicorefuzz import configspec
from unicorefuzz.configspec import serialize_spec, UNICOREFUZZ_SPEC
from unicorefuzz.harness import Harness
from unicorefuzz.unicorefuzz import Unicorefuzz


def getenv_default(envname: str, default: str) -> str:
    """
    Returns the environment variable if set, else returns the default
    :param envname: name of env variable to get
    :param default: what to return if envname is not set
    :return env variable or default if not set
    """
    env = os.getenv(envname)
    return env if env is not None else default


def load_conf(args: argparse.Namespace, silent: bool = False) -> Any:
    """
    Loads the config from args
    :param args: the arguments
    :param silent: If progress and config infos should be printed or not (if silent is True)
    :return: a loaded config
    """
    return configspec.load_config(args.config, silent)


# Note: The docstring will be used as commandline help for these funcs
def print_spec(args: argparse.Namespace) -> None:
    """
    Outputs expected config.py spec.
    :param args: the arguments
    """
    print(serialize_spec(UNICOREFUZZ_SPEC))


def wrap_probe(args: argparse.Namespace) -> None:
    """
    Attach, break and forward memory from target
    Former probewrapper.py
    """
    from unicorefuzz.probe_wrapper import ProbeWrapper

    ProbeWrapper(load_conf(args)).wrap_gdb_target()


def emulate(args: argparse.Namespace) -> None:
    """
    Drop the memory in the harness and start the emulation
    Former harness.py
    """
    Harness(load_conf(args)).harness(
        args.input_file, debug=args.debug, trace=args.trace, wait=args.wait
    )


def run_angr(args: argparse.Namespace) -> None:
    """
    Drop the memory in the angr harness and start concolic execution
    Former angr-harness.py
    """
    from unicorefuzz.angr_harness import AngrHarness

    AngrHarness(load_conf(args)).get_angry(args.input_file)


def wait_for_wrapper(args: argparse.Namespace, ucf: Unicorefuzz = None) -> None:
    """
    Blocks until data from probe wrapper becomes available
    """
    if ucf is not None:
        config = ucf.config
    else:
        config = load_conf(args)
    print("[*] Awaiting wrapper...")
    Unicorefuzz(config).wait_for_probe_wrapper()


def print_afl_path(args: argparse.Namespace) -> None:
    """
    print(Unicorefuzz(load_conf(args)).afl_path)
    """
    print(Unicorefuzz(load_conf(args, silent=True)).afl_path)


def fuzz(args: argparse.Namespace) -> None:
    """
    Starts afl using ucf emu
    """
    id = args.id
    restart = args.restart
    if restart and id != "0":
        raise ValueError("Only master (`id 0`) may `reset` the state.")
    if id == "0":
        id = "master"
        mode = "-M master"
    else:
        id = "fuzzer{}".format(id)
        mode = "-S {}".format(id)
    ucf = Unicorefuzz(load_conf(args))

    if restart:
        try:
            os.unlink(ucf.config.AFL_OUTPUTS)
        except:
            pass

    afl_inputs = ucf.config.AFL_INPUTS
    # See if output is already created, if not, we want to rerun afl instead of restart.
    if os.path.isdir(os.path.abspath(os.path.join(ucf.config.AFL_OUTPUTS, id))):
        print("[*] AFL path for node {} already exists. Resuming fuzzing.".format(id))
        afl_inputs = "-"

    wait_for_wrapper(args, ucf)

    config_path = ucf.config.path

    # Path to AFL: Should usually(tm) point to the AFLplusplus subrepo
    afl_path = getenv_default("AFL_PATH", ucf.afl_path)
    # Libunicorn_path: Unicorn allows us to switch out the native lib without reinstalling its python bindings.
    libunicorn_path = getenv_default("LIBUNICORN_PATH", ucf.libunicorn_path)
    # AFL_COMPCONV_LEVEL=2 is an awesome addition to afl-unicorn, and definitely the one you want :)
    # See afl++ repo for further infos
    afl_compcov_level = getenv_default("AFL_COMPCOV_LEVEL", "2")

    # TODO: forward all additional parameters to AFL directly, instead.
    afl_timeout = getenv_default("AFL_TIMEOUT", "4000+")

    emu_params = ""
    if args.trace:
        if not args.print_outputs:
            raise ValueError(
                "Won't accept debug option -t without -P. Slowdown without benefit."
            )
        emu_params += "-t "

    afl = os.path.join(afl_path, "afl-fuzz")
    ucf_main = os.path.join(ucf.config.UNICORE_PATH, "ucf")

    if os.getenv("UCF_DEBUG_START_GDB"):
        print("[d] UCF_DEBUG_START_GDB set. Starting GDB, raising AFL timeouts.")
        afl_timeout = "99999999+"
        afl = "{gdb} {afl} --args {afl}".format(gdb=ucf.config.GDB_PATH, afl=afl)

    env = 'PATH="{}:$PATH" LIBUNICORN_PATH="{}" AFL_COMPCOV_LEVEL="{}"'.format(
        afl_path, libunicorn_path, afl_compcov_level
    )
    if args.print_outputs:
        env = "{} AFL_DEBUG_CHILD_OUTPUT=1 ".format(env)
    if args.yolo:
        env = "{} AFL_NO_AFFINITY=1 AFL_SKIP_CPUFREQ=1 AFL_I_DONT_CARE_ABOUT_MISSING_CRASHES=1 ".format(
            env
        )

    run = (
        "{env} {afl} -U -m none -i {afl_in} -o {afl_out} -t {afl_timeout} {mode} "
        "-- python3 {ucf_main} emu {emu_params} -c {config_path} @@ || exit 1".format(
            env=env,
            afl=afl,
            afl_in=afl_inputs,
            afl_out=ucf.config.AFL_OUTPUTS,
            afl_timeout=afl_timeout,
            mode=mode,
            id=id,
            ucf_main=ucf_main,
            emu_params=emu_params,
            config_path=config_path,
        )
    )
    if args.print_outputs:
        print("[*] Starting: ", run)
    if os.getenv("UCF_DEBUG_PRINT_COMMAND_ONLY"):
        print("[d] ucf: Would execute:\n")
        print(run)
    else:
        os.system(run)


def kernel_setup(args: argparse.Namespace) -> None:
    """
    Sets up the kernel options needed to run AFL.
    """
    ucf = Unicorefuzz(load_conf(args))
    os.chdir(ucf.afl_path)
    run = "./afl-system-config || exit 1"
    if args.sudo:
        run = "sudo " + run
    print("[*] Setting AFL system conf")
    os.system(run)


# noinspection PyProtectedMember,PyDefaultArgument
def create_subparser(
    subparsers: argparse._SubParsersAction,
    name: str,
    func: Callable,
    aliases: Iterable[str] = [],
    uses_config: bool = True,
    uses_input: bool = False,
) -> argparse.ArgumentParser:
    """
    Creates and inits a subparser, initializing help from docstring
    :param subparsers: the initialized parser.add_subparsers
    :param name: the name for the new subparser
    :param func: the func to call (and to get the docstring from as help)
    :param aliases: set of aliases (other names), if needed
    :param uses_config: if the ucf config can be supplied using `-c`
    :param uses_input: if an input file can be supplied using `-i`
    :return: the initialized and added new subparser
    """
    subparser = subparsers.add_parser(name, aliases=aliases, help=func.__doc__)
    subparser.set_defaults(function=func)
    if uses_input:
        subparser.add_argument(
            "input_file",
            type=str,
            help="Path to the file containing the mutated input to load",
        )
    if uses_config:
        subparser.add_argument(
            "-c",
            "--config",
            type=str,
            default="config.py",
            help="The config file to use.",
        )
    return subparser


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Unicorefuzz, emulate kernels with AFL++-Unicorn"
    )
    subparsers = parser.add_subparsers(help="What unicorefuzz function to use.")
    probe_wrapper = create_subparser(subparsers, "attach", wrap_probe)

    harness = create_subparser(
        subparsers, "emu", emulate, aliases=["emulate"], uses_input=True
    )
    harness.add_argument(
        "-d",
        "--debug",
        default=False,
        action="store_true",
        help="Starts the testcase in uUdbg (if installed)",
    )
    harness.add_argument(
        "-t",
        "--trace",
        default=False,
        action="store_true",
        help="Enables debug tracing",
    )
    harness.add_argument(
        "-w",
        "--wait",
        default=True,
        action="store_true",
        help="Wait for the state directory to be present",
    )

    sub_fuzz = create_subparser(subparsers, "fuzz", fuzz)
    sub_fuzz.add_argument(
        "-i",
        "--id",
        type=str,
        default="0",
        help="The AFL multi fuzzer id to use (0 for master).",
    )
    sub_fuzz.add_argument(
        "-r",
        "--restart",
        default=False,
        action="store_false",
        help="If set, clears the afl_output directory before running.\nOnly works for master.\nDANGEROUS!!",
    )
    sub_fuzz.add_argument(
        "-P",
        "--print-outputs",
        default=False,
        action="store_true",
        help="When fuzing, print all child output (for debug)",
    )
    sub_fuzz.add_argument(
        "-t",
        "--trace",
        default=False,
        action="store_true",
        help="Enables debug tracing for children. Slow. Only useful with -P.",
    )
    sub_fuzz.add_argument(
        "-y",
        "--yolo",
        default=False,
        action="store_true",
        help="Ignore OS settings for coredump notifications and governor. Prefer to run `ucf sysconf -S` and not set this flag.",
    )

    sub_await = create_subparser(subparsers, "await", wait_for_wrapper)
    sub_afl_path = create_subparser(subparsers, "afl-path", print_afl_path)
    sub_spec = create_subparser(subparsers, "spec", print_spec)

    # Not yet ready for prime time :(
    # angr = create_subparser(
    #    subparsers, "angr", run_angr, aliases=["concolic"], uses_input=True
    # )

    init_system = create_subparser(subparsers, "sysconf", kernel_setup)
    init_system.add_argument(
        "-S",
        "--sudo",
        default=False,
        action="store_true",
        help="Auto escalate privileges",
    )

    args = parser.parse_args()

    if hasattr(args, "function"):
        args.function(args)
    else:
        parser.print_help()