VDB_ROOT = "/xxx/vdb_20120425"

import sys
sys.path.append(VDB_ROOT)

import vtrace
import threading
import traceback
import vdb
import time
import Elf
import os
from envi.archs.i386 import *

cmd = "./drmless %s" % sys.argv[1]

outputfile = sys.argv[2]

class MyBreakpoint(vtrace.Breakpoint):
    def __init__(self, address, callback):
        vtrace.Breakpoint.__init__(self, address)
        self.address = address
        self._cb = callback

    def notify(self, event, trace):
        self._cb(trace)

#
# Helper stuff
################

def put_in_file(fname,content):
    fd=open(fname,'w')
    fd.write(content)
    fd.close()

def ret1(trace):
  eax = trace.getRegister(REG_EAX)
  trace.setRegister(REG_EAX,1)

def ret0(trace):
  eax = trace.getRegister(REG_EAX)
  trace.setRegister(REG_EAX,0)  

import struct
dword = lambda x: struct.pack("<I",x)
undword = lambda x: struct.unpack("<I",x)[0]

def dump_to_file(trace):
  f = open(outputfile,'ab')
  esp = trace.getRegister(REG_ESP)
  ptr = undword(trace.readMemory(esp,4))
  f.write(trace.readMemory(ptr,16))
  f.close()

# MAIN
########

#Remove output file if it exists...
os.system("rm %s " % outputfile)

#Create empty drmlicense
orig_license = open('.drmlicense_orig').read()
zero_license = "\x00"*16

put_in_file(".drmlicense",zero_license)
trace = vtrace.getTrace()

trace.execute(cmd)
trace.setMode("RunForever", True)

#Say it's always drmprotected :)
trace.addBreakpoint(MyBreakpoint(0x0804EC2D, ret1)) 
trace.addBreakpoint(MyBreakpoint(0x0804B469, dump_to_file))
trace.run()