#!/usr/bin/env python
'''
Matthew Wollenweber
mjw@cyberwart.com
Wait for x iterations of function y

'''
DESC = "Skip through x iterations of function y"
USAGE = "!wait function_addr count"


import immlib
from immlib import LogBpHook
#import time
import struct
import unicodedata
import getopt


def main(args): 
    imm = immlib.Debugger()
    load_hook = MyLoadHook()

    addr = long(args[0], 16)
    count = int(args[1])
    
    imm.addKnowledge("wait_count", count)
    imm.addKnowledge("curr_count", 0)
    imm.addKnowledge("wait_bp", addr)
    
    load_hook.add("waiting hook", addr, 0, 0, 0)
    imm.Log("waiting hook set")
    
    return "waiting hook set"


def usage():
    print USAGE


if __name__=="__main__":
    print "This module is for use within Immunity Debugger only"


    
class MyLoadHook(LogBpHook):
    def __init__(self):
        LogBpHook.__init__(self)
        self.imm = immlib.Debugger()
             
    def run(self,regs):       
        imm = self.imm
        curr_count = imm.getKnowledge("curr_count")
        wait_count = imm.getKnowledge("wait_count")
        #imm.Log("curr count %i"  % (curr_count))
        if curr_count == wait_count:
            imm.setBreakpoint(long(imm.getKnowledge("wait_bp")))
            imm.Pause()
        else:
            curr_count += 1
            imm.forgetKnowledge("curr_count")
            imm.addKnowledge("curr_count", curr_count)
  
        return