/*
    Win32 Call Chains - Xref Mapper
    Pedram Amini

    This script is written for the purpose of mapping out the win32 API call
    chain. It will parse through an entire database and for each function output
    the call tree.

    This script is a dirty, disgusting, filthy little hack. I was in a rush to
    get the job done.
*/

#include <idc.idc>

// control debug output.
#define DEBUG

// control whether or not chainless routines are recorded in the output file.
#undef IGNORE_CHAINLESS

// control whether or not to exit when script is done.
#define EXIT_WHEN_DONE

////////////////////////////////////////////////////////////////////////////////
// main()
//
// called when script is first loaded.
//
// arguments: none.
// returns:   none.
//

static main ()
{
    auto called_func_name, called_func_raw_name, ea, first_func, func, func_name, idx;
    auto xref;
    auto outfile, fhandle;
    auto log_entry, num_calls;

    Message("[*] Pedram Amini - Call Mapper\n");
    Message("[*] Dumping each functions call list\n");

    // ask the user to select an output file.
    //if ((outfile = AskFile(1, "*.txt", "Select Output File")) == 0)
    //{
    //    Message("[!] No output file specified, quiting\n");
    //    return;
    //}

    // generate the output file name.
    outfile = GetInputFilePath() + ".calls";

    // try to open the specified file.
    if ((fhandle = fopen(outfile, "w+")) == 0)
    {
        Message("[!] Failed to open outfile: %s\n", outfile);
        return;
    }

    // we're about to start working.
    // change the status, make an announcement and disable dialogs.
    SetStatus(IDA_STATUS_WORK);
    Message("[-] Scanning through database ... \n");
    Batch(1);

    first_func = NextFunction(MinEA());

    // step through each function in the database.
    for (func = first_func; func != BADADDR; func = NextFunction(func))
    {
        // if the current function has no name, then move on to the next one.
        if (!FuncHasName(func))
            continue;

        // if the function is small, determine if it is a simple IAT wrapper.
        // XXX - this is by no means a perfect check.
        if ((GetFunctionAttr(func, FUNCATTR_END) - GetFunctionAttr(func, FUNCATTR_START)) <= 8)
        {
            // check for uncondition jump into data.
            if (GetMnem(func) == "jmp" && Dfirst(func) != BADADDR)
                continue;
        }

        // retrieve the function name.
        func_name = Demangle(GetFunctionName(func), INF_SHORT_DN);

        // if the called function doesn't have a demangled name, grab the regular name.
        if (func_name == 0)
            func_name = GetFunctionName(func);

        // if we still don't have a name, then move on.
        if (func_name == "")
            continue;

        // ignore function names containing:
        //      - single quote
        //      - question mark
        //      - at symbol
        if (strstr(func_name, "\'") != -1)
            continue;

        if (strstr(func_name, "?") != -1)
            continue;

        if (strstr(func_name, "@") != -1)
            continue;

        // ignore c++ operators (new/delete/etc).
        if (substr(func_name, 0, 9) == "operator ")
            continue;

        // convert Zw??? function names (from NTDLL.DLL) to Nt???
        if (substr(func_name, 0, 2) == "Zw")
            func_name = "Nt" + substr(func_name, 2, -1);

        #ifdef DEBUG
            Message("[d] Analyzing %s\n", func_name);
        #endif

        // create an array to store the data for this functions log entry.
        CreateArray("log_entry");
        log_entry = GetArrayId("log_entry");
        num_calls = 0;

        // step through each instruction in this function.
        for (ea = func; ea != FindFuncEnd(func); ea = NextNotTail(ea))
        {
            // ensure we are dealing with a call instruction.
            if (GetMnem(ea) != "call")
                continue;

            // determine the cross-referenced address.
            xref = Rfirst0(ea);

            // grab the function name.
            if (xref != BADADDR)
            {
                // if the function has a generic name (sub/loc/off_xxx), ignore it.
                if (!FuncHasName(xref))
                    continue;

                called_func_raw_name = GetFunctionName(xref);
            }
            // if we fail, try a data-reference.
            else
            {
                xref = Dfirst(ea);
            }

            // grab the function name.
            if (xref != BADADDR)
            {
                called_func_raw_name = Name(xref);
            }
            // if we still don't have a xref, move on.
            else
            {
                continue;
            }

            // retrieve the function name.
            called_func_name = Demangle(called_func_raw_name, INF_SHORT_DN);

            // if the called function doesn't have a demangled name, use the regular name.
            if (called_func_name == 0)
                called_func_name = called_func_raw_name;

            // if we still don't have a name, then move on.
            if (called_func_name == "")
                continue;

            // ignore function names containing:
            //      - single quote
            //      - question mark
            //      - at symbol
            if (strstr(called_func_name, "\'") != -1)
                continue;

            if (strstr(called_func_name, "?") != -1)
                continue;

            if (strstr(called_func_name, "@") != -1)
                continue;

            // ignore c++ operators (new/delete/etc).
            if (substr(called_func_name, 0, 9) == "operator ")
                continue;

            // convert Zw??? function names (from NTDLL.DLL) to Nt???
            if (substr(called_func_name, 0, 2) == "Zw")
                called_func_name = "Nt" + substr(called_func_name, 2, -1);

            // add the target of this call to the list.
            if (AddFunc(num_calls, called_func_name))
                num_calls++;
        }

        #ifdef IGNORE_CHAINLESS
            // ignore the routine if there are no intra-modular xrefs.
            if (num_calls == 0)
                continue;
        #endif

        // output the function name to the file.
        fprintf(fhandle, "%s" func_name);

        // step through the array and log the called functions list to the file.
        for (idx = 0; idx < num_calls; idx++)
        {
            #ifdef DEBUG
                Message(" -> %s", GetArrayElement(AR_STR, log_entry, idx));
            #endif

            fprintf(fhandle, ",%s", GetArrayElement(AR_STR, log_entry, idx));
        }

        #ifdef DEBUG
            Message("\n");
        #endif

        fprintf(fhandle, "\n");
        DeleteArray(log_entry);
    }

    fclose(fhandle);

    // we're done working.
    // change the status, make an announcement and enable dialogs.
    SetStatus(IDA_STATUS_READY);
    Message("[-] done. \n");
    Batch(0);

    // exit if necessary.
    #ifdef EXIT_WHEN_DONE
        Exit(0);
    #endif
}


////////////////////////////////////////////////////////////////////////////////
// AddFunc()
//
// add a function to the list.
//
// arguments: num_calls - index into our array to add the call.
//            func_name - function name to add.
// returns:   true if a function was added, false otherwise.
//

static AddFunc (num_calls, func_name)
{
    auto idx, log_entry;

    // grab a handle to our array.
    log_entry = GetArrayId("log_entry");

    // search for a duplicate entry.
    for (idx = 0; idx < num_calls; idx++)
    {
        // if a duplicate entry was found, then return immediately.
        if (func_name == GetArrayElement(AR_STR, log_entry, idx))
            return 0;
    }

    // add the function to our list.
    SetArrayString(log_entry, idx, func_name);

    return 1;
}


////////////////////////////////////////////////////////////////////////////////
// FuncHasName()
//
// check if a function has a real name or not.
//
// arguments: func_name - function name.
// returns:   boolean.
//

static FuncHasName (ea)
{
    if ((GetFlags(ea) & FF_LABL) == 0)
        return 1;

    return 0;
}