#!/usr/bin/env python

import sys, os, struct, string, re

# This should point to paimei if its not in site-packages
sys.path.append("C:\\code\\python\\paimei")

from pydbg import *
from pydbg.defines import *

######################################################################
#
# Config Class
#
######################################################################
class Config:
    def __init__(self, filename):
        self.filename = filename
        self.config = []
        
        self.parse_config()
        
    def parse_config(self):
        try:
            self.fh = open(self.filename, 'r')
        except:
            print "[!] Problem opening %s" % self.filename
        
        for line in self.fh.readlines():
            if not re.search("^#", line) and not re.search("^$", line):
                options = {}
                options["address"], options["type"], options["size"], options["data"] = line.split(",")
                options["address"] = string.atol(options["address"].strip(), 16)
                options["type"] = options["type"].strip()
                options["size"] = int(options["size"].strip())
                options["data"] = options["data"].strip()

                self.config.append(options)

    def get_config(self):
        return self.config

######################################################################
#
# Entries Class
#
######################################################################

class Entries:
    def __init__(self, config, dbg):
        self.entries = {}

        for item in config:
            if self.entries.has_key(item["address"]):
                self.entries[item["address"]].append(item)
            else:
                self.entries[item["address"]] = []
                self.entries[item["address"]].append(item)
                
                print "[*] Setting bp @ 0x%08x" % item["address"]
                dbg.bp_set(item["address"], restore=False, handler=self.class_bp_handler)
        
    def class_bp_handler(self, dbg):
        address = dbg.exception_address
        delayed_entries = []
        
        if not self.entries.has_key(address):
            return DBG_CONTINUE
        
        entry = self.entries[address]
        for item in entry:
            if item["data"] == "1":
                delayed_entries.append(item)
            else:
                deref = DerefData(item, dbg)
            
        for item in delayed_entries:
            deref = DerefData(item, dbg)
        
        return DBG_CONTINUE

    def debug_dump_item(self, item):
        print "0x%08x,%c,%d,%s" % (item["address"], item["type"], item["size"], item["data"])
        
        return True

######################################################################
#
# DerefData Class
#
######################################################################

class DerefData:
    def __init__(self, config, dbg):
        self.address = config["address"]
        self.type = config["type"]
        self.size = config["size"]
        self.data = config["data"]
        self.deref_data = ""
        
        #print "Got %x %c %s" % (self.address, self.type, self.data)
        if self.data != "1":
            self.deref(dbg)
        else:
            dbg.set_callback(EXCEPTION_SINGLE_STEP, self.class_ss_handler)
            dbg.single_step(True)

    def class_ss_handler(self, dbg):
        self.deref(dbg)
        dbg.single_step(False)
        
        return DBG_CONTINUE
    
    #
    # Our deref function whick branches based on type to appropriate handler
    # 
    # Not really excited about the output style but its *decent*
    #
    def deref(self, dbg):
        if self.type == 'r':
            self.deref_data = self.get_register(dbg)
            print "[*] 0x%08x %8s [%-7s] is 0x%-8x [%d]" % (self.address, self.data, "Reg", self.deref_data, self.size)         
        elif self.type == 'g':
            self.deref_data = self.get_global(dbg)
            print "[*] %-8s 0x%08x is 0x%x" % ("Global", self.address, self.deref_data)
        elif self.type == 'l':
            self.deref_data = self.get_local(dbg)
            print "[*] %-8s 0x%08x is 0x%x" % ("Local", self.address, self.deref_data)
        elif self.type == 'o':
            self.deref_data = self.get_offset(dbg)
            print "[*] 0x%08x %8s [%-7s] is 0x%-8x [%d]" % (self.address, self.data, "Offset", self.deref_data, self.size)
        elif self.type == 'p':
            self.deref_data = self.get_pointer(dbg)
            print "[*] 0x%08x %8s [%-7s] is 0x%-8x [%d]" % (self.address, self.data, "Pointer", self.deref_data, self.size)
        else:
            print "[!] Unknown type %s" % (self.type)
            return False
            
        return True
    
    #
    # Returns the value of a register (no dereferencing is performed)
    #
    def get_register(self, dbg):
        if self.data > "3":
            return self.get_reg_value(dbg, self.data)
        else:
            op = self.get_op(dbg)
            return self.get_reg_value(dbg, op.reg)
    
        return False
    
    #
    # Global is technically a pointer but i left this
    #
    def get_global(self, dbg):
        return self.flip_value(dbg.read_process_memory(self.address, self.size))
    
    #
    # Local is technically an offset but i left this...
    #
    def get_local(self, dbg):
        op = self.get_op(dbg)
        reg = self.get_reg_value(dbg, op.basereg)
        off = op.displacement
        addr = reg + off

        return self.flip_value(dbg.read_process_memory(addr, self.size))
    
    #
    # Returns the dereferenced value after calculating the address (basereg+displacement)
    #
    def get_offset(self, dbg):
        op = self.get_op(dbg)
        reg = self.get_reg_value(dbg, op.basereg)
        off = op.displacement
        addr = reg + off
        
        return self.flip_value(dbg.read_process_memory(addr, self.size))

    #
    # Returns the value of pointer after it has been dereferenced
    #
    def get_pointer(self, dbg):
        if self.data > "3":
            return self.flip_value(dbg.read_process_memory(string.atol(self.data, 16), self.size))
        else:    
            op = self.get_op(dbg)
            reg = self.get_reg_value(dbg, op.basereg)
            
            return self.flip_value(dbg.read_process_memory(reg, self.size))
        
        return False
    
    #
    # Returns the proper value in little endian
    #   
    def flip_value(self, value):
        if self.size == 1:
            return struct.unpack("<B", value)[0]
        elif self.size == 2:
            return struct.unpack("<H", value)[0]
        elif self.size == 4:
            return struct.unpack("<L", value)[0]
        elif self.size == 8:
            return struct.unpack("<Q", value)[0]
        else:
            return False
        
        return False
    
    #
    # Returns a pydasm operand class
    #
    def get_op(self, dbg):
        if self.data == "1":
            return dbg.get_instruction(self.address).op1
        if self.data == "2":
            return dbg.get_instruction(self.address).op2
        elif self.data == "3":
            return dbg.get_instruction(self.address).op3
        else:
            return False
        
        return False

    #
    # Returns a registers contents
    #
    def get_reg_value(self, 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
        elif register == "AX": return context.Eax & 0x0000FFFF
        elif register == "CX": return context.Ecx & 0x0000FFFF
        elif register == "DX": return context.Edx & 0x0000FFFF
        elif register == "BX": return context.Ebx & 0x0000FFFF
        elif register == "SI": return context.Esi & 0x0000FFFF
        elif register == "DI": return context.Edi & 0x0000FFFF
        elif register == "AH" or register == 20: return ((context.Eax & 0xFF00) >> 8)
        elif register == "AL" or register == 16: return context.Eax & 0x00FF
        elif register == "CH" or register == 21: return ((context.Ecx & 0xFF00) >> 8)
        elif register == "CL" or register == 17: return context.Ecx & 0x00FF            
        elif register == "DH" or register == 22: return ((context.Edx & 0xFF00) >> 8)
        elif register == "DL" or register == 18: return context.Edx & 0x00FF
        elif register == "BH" or register == 23: return ((context.Ebx & 0xFF00) >> 8)
        elif register == "BL" or register == 19: return context.Ebx & 0x00FF
        else: return False
            
        return False
    
    #
    # Debug dump of class variables
    #    
    def debug_dump(self):
        print "Dump 0x%08x,%c,%d,%s" % (self.address, self.type, self.size, self.data)
        
        return True
    
    def debug_disasm(self, op):
        print "[*] Type 0x%08x" % op.type
        print "[*] Reg 0x%08x" % op.reg
        print "[*] BaseReg 0x%08x" % op.basereg
        print "[*] IndexReg 0x%08x" % op.indexreg
        print "[*] Scale 0x%08x" % op.scale
        print "[*] DispBytes 0x%08x" % op.dispbytes
        print "[*] DispOffset 0x%08x" % op.dispoffset
        print "[*] ImmBytes 0x%08x" % op.immbytes
        print "[*] ImmOffset 0x%08x" % op.immoffset
        print "[*] SectionBytes 0x%08x" % op.sectionbytes
        print "[*] Section 0x%08x" % op.section
        print "[*] Displacement 0x%08x" % op.displacement
        print "[*] Immediate 0x%08x" % op.immediate
        print "[*] Flags 0x%08x" % op.flags
    
        return True
                
######################################################################
#
# Breakpoint handlers
#
######################################################################

def handler_breakpoint(dbg):
    # Initial module bp we need to process entries
    if dbg.first_breakpoint:
        entries = Entries(dbg.args, dbg)

        return DBG_CONTINUE
    
    print "[!] Unknown bp caught @ 0%08x" % dbg.exception_address
    
    return DBG_CONTINUE

######################################################################
#
# Helper functions
#
######################################################################

#
# Ghetto idc output to import into ida if wanted...not really thought out or tested
#
def output_idc(filename, comments):
    fh = open(filename, "w")
    
    fh.write("#include <idc.idc>\n\n")
    fh.write("static main()\n")
    fh.write("{\n")
    fh.write("\tauto previous;\n\n")
    
    for comment in comments:
        fh.write("\tprevious = Comment(%x);\n" % comment["address"])
        fh.write("\tMakeComm(%x, previous + \"\\n\" + %s);\n" % (comment["address"], comment["comment"]))
    
    fh.write("}\n")
    fh.close()
    
    return True

#
# Attaches to procname if it finds it otherwise loads
#
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:
            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

######################################################################
#
# Command line arguments
#
######################################################################

if len(sys.argv) < 3:
    print "Usage: %s <process name> <la conf file>" % sys.argv[0]
    
    sys.exit(-1)

procname = sys.argv[1]
configfile = sys.argv[2]

config = Config(configfile).get_config()

dbg = pydbg()
dbg.procname = procname
dbg.args = config

dbg.set_callback(EXCEPTION_BREAKPOINT, handler_breakpoint)

if not attach_target_proc(dbg, procname):
    print "[!] Couldnt load/attach to %s" % procname
    
    sys.exit(-1)

dbg.debug_event_loop()