#!/usr/bin/env python3
import logging

import click
import rzpipe


@click.command()
@click.option("--exe_path", help="Path to the map file", required=True)
@click.option("--map_path", help="Path to the map file", required=True)
@click.option("--project_name", default="analysed_bin", help="Name of the project (default: analysed_bin)", required=True)
@click.option("--rizin_cmd", default="aaaa", help="Rizin command to run to analyse the bin (default: aaaa)", required=True)
@click.option("--debug", is_flag=True, help="Enable debug logging", required=False)
def analyse_binary(exe_path, map_path, rizin_cmd, project_name, debug):
    log_level = logging.DEBUG if debug else logging.INFO
    logging.basicConfig(level=log_level,
                        format='%(asctime)s - %(levelname)s - %(message)s',
                        datefmt='%Y-%m-%d %H:%M:%S')

    rz = rzpipe.open(exe_path)

    logging.debug("screaming in the void...")
    rz.cmd(rizin_cmd)
    logging.debug("done screaming")

    map_entries = {}
    with open(map_path, "r") as map_file:
        for line in map_file:
            _, _, func_and_offset = line.split(" ")
            func_name, offset = func_and_offset.rsplit("_", 1)
            rhex = "0x{}".format(offset.strip())
            name = func_name.strip()
            if rhex and name:
                logging.debug("Adding {} - {} to map".format(rhex, name))
                map_entries[rhex] = name

    logging.debug("Map done, running aflj")
    functions = rz.cmdj("aflj")
    for function in functions:
        offset_formated = "0x{:08x}".format(function['offset'])
        if offset_formated in map_entries:
            new_name = map_entries[offset_formated]
            logging.debug("Renaming {} to {}".format(offset_formated, new_name))
            rz.cmd("s {}".format(offset_formated))
            rz.cmd("afn {}".format(new_name))

    rz.cmd("Ps {}".format(project_name))
    logging.info("Project {} created".format(project_name))


if __name__ == "__main__":
    analyse_binary()