import argparse
import os
import subprocess
import time
from pathlib import Path

import frida  # type: ignore
from androguard.core.apk import get_apkid  # type: ignore

FRIDA_SCRIPT = Path(__file__).parent / "hook.js"
STACK_CONSUMER_B64 = Path(__file__).parent / "StackConsumer.dex.b64"


# Define handler to event generated by the scripts
def on_message(message, data):
    if message["type"] == "error":
        print(f"[error] {message['description']}")
        print(message["stack"])
    elif message["type"] == "send" and message["payload"]["type"] == "invoke":
        handle_invoke_data(message["payload"]["data"])
    elif message["type"] == "send" and message["payload"]["type"] == "class-new-inst":
        handle_class_new_inst_data(message["payload"]["data"])
    elif message["type"] == "send" and message["payload"]["type"] == "cnstr-new-isnt":
        handle_cnstr_new_inst_data(message["payload"]["data"])
    else:
        print("[on_message] message:", message)


def print_stack(stack, prefix: str):
    for frame in stack:
        native = ""
        if frame["is_native"]:
            native = " (native)"
        print(f"{prefix}{frame['method']}:{frame['bytecode_index']}{native}")


# def get_ty(java_name: str) -> str:
#    """Return the android name from the java name of a class / type"""
#    # TODO: array
#    # TODO: scalar
#    if java_name == "V":  # tmp stub
#        return "V"
#    return f"L{java_name.replace('.', '/')};"


# def get_method_id(method_data) -> str:
#    """Get a method descriptor from the different elements collected from the methods."""
#    name = method_data["name"]
#    ret = get_ty(method_data["ret"])
#    cls = get_ty(method_data["class"])
#    args = "".join(map(get_ty, method_data["args"]))
#    return f"{cls}->{name}({args}){ret}"


def handle_invoke_data(data):
    method = data["method"]
    # caller_method = "?"  # get_method_id(data["caller_method"])
    # addr = data["addr"]
    print("Method.Invoke:")
    print(f"    called: {method}")
    print(f"    stack:")
    print_stack(data["stack"], "        ")
    # print(f"    by:     {caller_method}")
    # print(f"    at:     0x{addr:08x}")


def handle_class_new_inst_data(data):
    constructor = data["constructor"]
    # caller_method = "?"  # get_method_id(data["caller_method"])
    # addr = data["addr"]
    print("Class.NewInstance:")
    print(f"    called: {constructor}")
    print(f"    stack:")
    print_stack(data["stack"], "        ")
    # print(f"    by:     {caller_method}")
    # print(f"    at:     0x{addr:08x}")


def handle_cnstr_new_inst_data(data):
    constructor = data["constructor"]
    if not constructor.startswith("Lcom/example/theseus"):
        return
    # caller_method = "?"  # get_method_id(data["caller_method"])
    # addr = data["addr"]
    print("Constructor.newInstance:")
    print(f"    called: {constructor}")
    print(f"    stack:")
    print_stack(data["stack"], "        ")
    # print(f"    by:     {caller_method}")
    # print(f"    at:     0x{addr:08x}")


def main():
    parser = argparse.ArgumentParser(
        prog="Android Theseus project",
    )
    parser.add_argument(
        "-a", "--apk", required=True, help="Target application", type=Path
    )
    parser.add_argument(
        "-s",
        "--device",
        default="",
        help="The android device to connect to, eg: 'emulator-5554'",
        type=str,
    )
    args = parser.parse_args()
    env = dict(os.environ)

    if args.device != "":
        device = frida.get_device(args.device)
        env["ANDROID_SERIAL"] = args.device
    else:
        device = frida.get_usb_device()

    app = get_apkid(args.apk)[0]

    if device.enumerate_applications([app]):
        # Uninstall the APK if it already exist
        subprocess.run(["adb", "uninstall", app], env=env)
    subprocess.run(["adb", "install", str(args.apk.absolute())], env=env)

    with FRIDA_SCRIPT.open("r") as file:
        script = file.read()
    with STACK_CONSUMER_B64.open("r") as file:
        script = script.replace(
            "<PYTHON REPLACE StackConsumer.dex.b64>",
            file.read().replace("\n", "").strip(),
        )

    pid = device.spawn([app])
    session = device.attach(pid)
    script = session.create_script(script)

    script.on(
        "message",
        on_message,
    )

    # Load script
    script.load()
    # Resume the execution of the APK
    device.resume(pid)

    print("Press ENTER to finish the analysis")
    input()