"""
ApiHash calculator script (c) Hex-Rays
"""
import re
import idaapi
from idaapi import Choose2
from ctypes import *

SANE_NAME_RE = re.compile("[@?$:`+&\[\]]", 0)

# -----------------------------------------------------------------------
def sanitize_name(name):
    return SANE_NAME_RE.sub("_", name)

# -----------------------------------------------------------------------
class create_cfunc_t(object):
    """
    Utility class to execute a sequence of bytes in the context
    of the Python interpreter.
    It allocates an executable memory page, copies the bytes and calls
    them. It uses ctypes to call the bytes.
    """
    def __init__(self, proto, body):
        # Alloc as: MEM_COMMIT, PAGE_EXECUTE_READWRITE
        func_ptr = windll.kernel32.VirtualAlloc(0, len(body), 0x1000, 0x40 )

        # Copy the body to the allocated space
        memmove(func_ptr, body, len(body))

        # Instantiate function
        self.func = proto(func_ptr)

    def __call__(self, name):
        return self.func(name, 0)

# -----------------------------------------------------------------------
def DelEnumByName(name):
    for i in range(0, idc.GetEnumQty()):
        id = idc.GetnEnum(i)
        if idc.GetEnumName(id) != name:
            continue
        idc.DelEnum(id)
        return True
    return False

# -----------------------------------------------------------------------
def extract_function_bytes(name):
    """Given a function name, this function returns the function's bytes"""
    f = idaapi.get_func(idc.LocByName(name))
    if f is None:
        return None
    return idaapi.get_many_bytes(f.startEA, f.endEA - f.startEA)
# (debug)
#    f = open('body.bin', 'rb')
#    b = f.read()
#    f.close()
    return b

# -----------------------------------------------------------------------
def fetch_debug_names():
    """Returns the debug names as a list of tuples(apiaddr, apiname, module_name)"""
    ret = []
    dn = idaapi.get_debug_names(idaapi.cvar.inf.minEA, idaapi.cvar.inf.maxEA)
    for addr in dn:
        n = dn[addr]
        i = n.find('_')
        # Append as (api addr, api name, module name)
        ret.append((addr, n[i+1:], n[:i]))

    return ret

# -----------------------------------------------------------------------
class ApiHashChoose(Choose2):
    """Api hash chooser class"""

    def __init__(self):
        Choose2.__init__(
            self,
            "API hash calculator",
            [ ["Module", 15], ["Hash", 8], ["API", 40] ]
        )

        # Populate hashes
        self.calc_hashes()
        self.icon = 5
        self.popup_names = ["Import as Enum", "", "Export to file", "Refresh"]


    def init_calc_hash(self):
        # Extract function body from the database
        body = extract_function_bytes('calc_hash')
        if not body:
            idc.Warning("Failed to extract function bytes. No function named 'calc_hash' was found!")
            return False

        # Create a function
        self.calc_hash = create_cfunc_t(
            # prototype: uint32 func(char *, uint32)
            WINFUNCTYPE(c_uint32, c_char_p, c_uint32),
            body)

        return True


    def calc_hashes(self):
        while True:
            if not self.init_calc_hash():
                break

            dn = fetch_debug_names()
            if not dn:
                idc.Warning("Failed to retrieve debug names. Make sure debugger is active!")
                break

            # Format the debug names
            cache = {}
            for addr, name, modname in dn:
                hash = self.calc_hash(name)
                if modname not in cache:
                    cache[modname] = []
                cache[modname].append((name, hash, addr))

            # Create items to be displayed
            items = []
            for modname in cache:
                mod = cache[modname]
                for name, hash, addr in mod:
                    items.append([modname, "%08X" % hash, name])

            self.items = items
            self.cache = cache

            return

        # No hashes
        self.items = [ ["n/a", "n/a", "n/a"] ]
        self.cache = {}


    def OnClose(self):
        global APIHASH_CHOOSE
        del APIHASH_CHOOSE


    def OnEditLine(self, n):
        # Export to a text file
        fn = idc.AskFile(1, "*.txt", "Export as")
        if not fn:
            return 1
        try:
            f = open(fn, 'w')
        except:
            idc.Warning('Could not create file: %s' % fn)
            return 1

        for modname in self.cache:
            funcs = self.cache[modname]
            f.write('Module: %s\n' % modname)
            for name, hash, _ in funcs:
                f.write('\t%08X -> %s\n' % (hash, name))

        f.close()


    def OnInsertLine(self):
        try:
            for modname in self.cache:
                t = sanitize_name(modname + "_hashes")
                DelEnumByName(t)
                id = idc.AddEnum(0, t, idaapi.hexflag())
                funcs = self.cache[modname]
                for name, hash, _ in funcs:
                    name = sanitize_name(name)
                    idc.AddConstEx(id, "hash_%s_%s" % (modname, name), hash, idaapi.BADADDR)
        except Exception, e:
            print "Failed while importing enums!\n" + str(e)


    def OnGetLine(self, n):
        return self.items[n]


    def OnGetSize(self):
        return len(self.items)


    def OnRefresh(self, n):
        self.calc_hashes()
        return n


    def show(self):
        return False if self.Show() < 0 else True

# -----------------------------------------------------------------------
try:
    APIHASH_CHOOSE
except:
    APIHASH_CHOOSE = ApiHashChoose()

APIHASH_CHOOSE.show()