from binaryninja import *
import os
from collections import defaultdict
from pipes import quote
from datetime import datetime
class_dump = quote('/Users/stefan/Downloads/class-dump')

def log_info(msg):
    log(1,"LOG - "+ str(datetime.now()) + "    " + msg)

def parse_objc(bv):
    log_info("Parsing obj-c metadata")
    file = bv.file.filename

    command = class_dump + ' -A ' + quote(file) + ' | grep -E "interface|//" | grep -v "@property" | sed -ne \'/@interface/,$p\' | less'
    parsed_functions = defaultdict(list)
    last_interface = None
    replacements = os.popen(command).read()
    for line in replacements.split("\n"):
        if line is None:
            continue

        interface_pos = line.find("@interface")
        colon_pos = line.find(":")
        comment_pos = line.find("//")
        method_pos = line.find("(")

        if interface_pos != -1:
            last_interface = line[11:colon_pos]

        if comment_pos > 0 and last_interface is not None:
            # We have a method for an interface
            method_type_pos = line.find("+")
            if method_type_pos != -1:
                method_type = "+"
            else:
                method_type = "-"

            method_name = line[method_pos:comment_pos].rstrip()
            method_name = method_name[:-1] + ":"  # replace semicolon with colon
            addr_pos = line.find("0x")
            address = line[addr_pos:]
            method_signature = method_type + '[' + last_interface + method_name + ']'

            parsed_functions[address] = method_signature

    log_info("Done parsing class-dump output. Replacing functions...")
    functions_replaced = 0
    for binja_func in bv.functions:
        for method_addr, method_signature in parsed_functions.iteritems():
            address = int(method_addr, 16)
            binja_addr = binja_func.symbol.address
            if binja_addr == address:
                # print "Replacing binja_func! " + binja_func.name + " -> " + method_signature
                functions_replaced += 1
                binja_func.name = method_signature

    log_info("Done replacing functions.")
    log_info(str(functions_replaced) + " functions replaced!")

PluginCommand.register("Parse Obj-c", "parse obj-c metadata", parse_objc)