/*********************************************************************
* vtable_xref.idc
*
* Copyright (c) 2008 Luis Miras
* Licensed under the BSD License
*
*********************************************************************/

#include <idc.idc>


/*********************************************************************
* Function: breakpointHandler
*
* This function is used as the expression for conditional breakpoints
* on vtable functions. The caller is calculated and a cross reference
* is created.
*********************************************************************/
static breakpointHandler()
{
    auto caller;
    caller = PrevHead(Dword(ESP), (Dword(ESP) - 10));
    AddCodeXref(caller,EIP , XREF_USER | fl_CN);
    return 0; // don't stop on breakpoint
}

/*********************************************************************
* Function: clearBPs
*
* This function operates on a selected vtable. Any breakpoints on 
* vtable functions are cleared.
*********************************************************************/
static clearBPs()
{
    auto currAddr;
    auto vStart;
    auto vEnd;
    auto virFunc;
    
    vStart = SelStart();
    vEnd = SelEnd();
    
    Message("[-] clearBPs() start: %x  end: %x\n", vStart, vEnd);

    if ((vStart == BADADDR) || (vEnd  == BADADDR))
    {
        Message("** No selection made!\n");
        return;    
    }
    
    if ((vStart - vEnd) %4 != 0)
    {
        Message("** not DWORD aligned!\n");
        return;    
    }
    
    for (currAddr = vStart; currAddr < vEnd; currAddr = currAddr + 4)
    {
        virFunc = Dword(currAddr);
       
        if (GetBptAttr(virFunc, BPTATTR_EA) != -1) // is bpt set?
        {
            if (DelBpt(virFunc) == 0)
            {
                Message("** %x DelBpt() failed!  %s\n;", virFunc, Name(virFunc));
                return;
            }
            else
            {
                Message("   %x clear          %s\n", virFunc, Name(virFunc));
            }
        }
        else
        {
            Message("   %x already clear  %s\n", virFunc, Name(virFunc));
        }
    }   
}


/*********************************************************************
* Function: setBPs
*
* This function operates on a selected vtable. Conditional breakpoints
* are set on all vtable functions. The conditional expression is set 
* to the breakpointHandler() function.
*********************************************************************/
static setBPs()
{
    auto currAddr;
    auto vStart;
    auto vEnd;
    auto virFunc;
    
    vStart = SelStart();
    vEnd = SelEnd();
    
    Message("[+] setBPs() start: %x  end: %x\n", vStart, vEnd);

    if ((vStart == BADADDR) || (vEnd  == BADADDR))
    {
        Message("** No selection made!\n");
        return;    
    }
    
    if ((vStart - vEnd) %4 != 0)
    {
        Message("** not DWORD aligned!\n");
        return;    
    }
    
    for (currAddr = vStart; currAddr < vEnd; currAddr = currAddr + 4)
    {
        virFunc = Dword(currAddr);
       
        if (GetBptAttr(virFunc, BPTATTR_EA) == -1) // no bpt 
        {
            if (!AddBptEx(virFunc, 0, BPT_SOFT))
            {
                Message("** %x AddBptEx() failed!  %s\n;", virFunc, Name(virFunc));
                return;
            }

            if (!SetBptCnd(virFunc, "breakpointHandler()"))
            {
                Message("** %x SetBptCnd() failed!  %s\n;", virFunc, Name(virFunc));
                return;
            }

            Message("   %x set            %s\n", virFunc, Name(virFunc));
        }
        else
        {
            Message("   %x already set    %s\n", virFunc, Name(virFunc));
        }
    }   
}

/*********************************************************************
* Function: main
*
* main's only purpose is to add hot keys.
*********************************************************************/
static main()
{
    AddHotkey("Alt-F7", "setBPs");
    AddHotkey("Alt-F8", "clearBPs");
}