from __future__ import print_function
from rich.console import Console
from rich.syntax import Syntax

console = Console()

import frida
import sys
import argparse
import json
import time 

parser = argparse.ArgumentParser(description='Process some integers.')
parser.add_argument('-s', '--spawn', action='store_true', help='spawn process')
parser.add_argument('-U', '--usb', action='store_true', help='use usb device')
parser.add_argument('-n', '--name', required=True, type=str, help='name of process')
parser.add_argument('methods', metavar='SEL', type=str, nargs='+',
                    help='a method selector like "*[NSMutable* initWith*]"')

args = parser.parse_args()

device = frida.get_local_device()
if args.usb:
    device = frida.get_usb_device()

name = args.name 
if name.isdigit():
    name = int(name)

if args.spawn:
    session = device.spawn(name)
else:
    session = device.attach(name)

print(f"Attached to {name} on {device}")

try:
    js = open("funtime.js").read()
    script = session.create_script(js)

    def format_call(info):
        obj = info["args"][0]
        parts = info["targetMethod"].split(":")
        objstr = f"""({obj["typeDescription"]})( {obj["objectDescription"]} )"""
        if info["targetType"] == "+":
            objstr = obj["typeDescription"].split(" ")[0]

        formatted = "\n\t" + f"""{info["targetType"]}[{objstr} """.replace("\n", "\n\t")

        if len(parts) == 1:
            formatted += "\n\t\t" + parts[0]
        else:
            for i, arg in enumerate(info["args"][2:]):
                formatted += ("\n\t\t" + f"""{parts[i]}: ({arg["typeDescription"]})""" + 
                    f"""( {arg["objectDescription"]} )""".replace("\n", "\n\t\t"))

        formatted += "];\n\t"

        if info["retTypeDescription"] != "void":
            formatted += (f"""return ({info["retTypeDescription"]})""" + 
                f"""( {info["returnDescription"]} );""".replace("\n", "\n\t\t"))

        return formatted + f" // time: {time.time()}"

    def on_message(message, data):
        if "payload" in message:
            payload = json.loads(message["payload"])
            formatted = format_call(payload)
            console.print(Syntax(formatted, "objc", background_color="black"))

    script.on('message', on_message)
    script.load()

    for method in args.methods:
        try:
            script.exports.hook(method)
        except Exception as e:
            print("error", e)

    sys.stdin.read()
except KeyboardInterrupt:
    pass

# script.unload()
session.detach()