'''
    IDA PyDbg Loader
    
    (c) Cody Pierce '08
    
    Description: A simple PyDbg launcher demonstrating access from IDA.
'''

import sys

# Specify our path to pydbg
sys.path.append(r'c:\code\python\paimei')

from pydbg import *
from pydbg.defines import *

# Helper Function
def attach_target_proc(dbg, procname):
    imagename = procname.rsplit('\\')[-1]
    print "[*] Trying to attach to existing %s" % imagename
    for (pid, name) in dbg.enumerate_processes():
        if imagename in name.lower():
            try:
                print "[*] Attaching to %s (%d)" % (name, pid)
                dbg.attach(pid)
            except:
                print "[!] Problem attaching to %s" % name
                
                return False
            
            return True
    
    try:
        print "[*] Trying to load %s" % (procname)
        dbg.load(procname, "")
        
    except:
        print "[!] Problem loading %s" % (procname)
        
        return False
    
    return True

# We need to set our bps from IDA here, then let pydbgc walk
def breakpoint_handler(dbg):
    # Initial module bp
    if dbg.first_breakpoint:
        # Set IDA BP
        dbg.bp_set(dbg.ida_ea, restore=False)
    
        return DBG_CONTINUE

    print_state(dbg)
    dbg.detach()
    
    return DBG_CONTINUE
    
def get_reg_value(dbg, register):
    context = dbg.get_thread_context(dbg.h_thread)

    if   register == "eax" or register == 0: return context.Eax
    elif register == "ecx" or register == 1: return context.Ecx
    elif register == "edx" or register == 2: return context.Edx
    elif register == "ebx" or register == 3: return context.Ebx
    elif register == "esp" or register == 4: return context.Esp
    elif register == "ebp" or register == 5: return context.Ebp
    elif register == "esi" or register == 6: return context.Esi
    elif register == "edi" or register == 7: return context.Edi
    elif register == "eip" or register == 8: return context.Eip

    return False
    
def print_state(dbg):
    address = get_reg_value(dbg, "eip")
    instruction = dbg.get_instruction(address)

    '''
    eax=7ffdf000 ebx=00000001 ecx=00000002 edx=00000003 esi=00000004 edi=00000005
    eip=7c901230 esp=0092ffcc ebp=0092fff4 iopl=0         nv up ei pl zr na pe nc
    cs=001b  ss=0023  ds=0023  es=0023  fs=0038  gs=0000             efl=00000246
    ntdll!DbgBreakPoint:
    7c901230 cc              int     3
    '''

    try:
        module = dbg.addr_to_module(address).szModule
    except:
        module = "N/A"

    sys.stdout.write("\n")
    sys.stdout.write("eax=%08x ebx=%08x ecx=%08x edx=%08x esi=%08x edi=%08x\n" %
    (get_reg_value(dbg, "eax"), get_reg_value(dbg, "ebx"), get_reg_value(dbg, "ecx"),
    get_reg_value(dbg, "edx"), get_reg_value(dbg, "esi"), get_reg_value(dbg, "edi")))
    sys.stdout.write("eip=%08x esp=%08x ebp=%08x\n\n" %
    (get_reg_value(dbg, "eip"), get_reg_value(dbg, "esp"), get_reg_value(dbg, "ebp")))
    sys.stdout.write("%s!%08x  %s\n\n" %
    (module, address, dbg.disasm(address)))

    return None
    
# IDC Functions
ea = ScreenEA()
process = GetInputFile()

# PyDbg Functions
dbg = pydbg()
dbg.ida_ea = ea
dbg.ida_process = process
dbg.set_callback(EXCEPTION_BREAKPOINT, breakpoint_handler)

if attach_target_proc(dbg, process):
    dbg.debug_event_loop()
else:
    print "[!] Couldnt load/attach to %s" % process