#Checkfault
#Author: St0rn
#Use python-ptrace
#
#Usage: checkfault.py [Type] [Binary] [Length]
#
#!/usr/bin/env python


import os
import sys
from ptrace.debugger.debugger import PtraceDebugger
from ptrace.debugger.child import createChild
from ptrace.tools import locateProgram

def generatejunk(length): 
 taba = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
 tabb = "abcdefghijklmnopqrstuvwxyz"
 tabc = "0123456789"
 
 junk = ""
 a = 0 
 b = 0
 c = 0
 
 while len(junk) < length:
  junk += taba[a] + tabb[b] + tabc[c]
  c += 1
  if c == len(tabc): 
   c = 0
   b += 1 
  if b == len(tabb):
   b = 0
   a += 1
  if c == len(tabc):
   a = 0
 return junk

#Find EIP control offset Function
def checkfault(binary, types, length):
 #types = arg, get
 def trace(program):
  env = None

  return createChild(program, False, env)
 
 if types == "arg":
  junk = generatejunk(length)
  payload = []
 
  payload.append(locateProgram(binary))
  payload.append(junk)
  pid = trace(payload)
  
  debug = PtraceDebugger()
  ps = debug.addProcess(pid, True)
  
  ps.cont()
  try: 
   signal = ps.waitSignals()
   print "\n%s Detected!" % signal
   addr_eip = hex(ps.getInstrPointer())
   addr_ebp = hex(ps.getFramePointer())

   ebp = addr_ebp[8] + addr_ebp[9] + addr_ebp[6] + addr_ebp[7] + addr_ebp[4] + addr_ebp[5] + addr_ebp[2] + addr_ebp[3]

   eip = addr_eip[8] + addr_eip[9] + addr_eip[6] + addr_eip[7] + addr_eip[4] + addr_eip[5] + addr_eip[2] + addr_eip[3]
  except:
   print "\nNo SIGSEGV detected!\nMaybe the program is not vulnerable\nOr try with more bytes\n"
   debug.quit()
   sys.exit()

  try:
   print "\nControl of EBP after " + str(junk.index(ebp.decode("hex"))) + " Bytes"
  except:
   print "\nNo Control of EBP"


  try:
   print "Control of EIP after " + str(junk.index(eip.decode("hex"))) + " Bytes\n"
  except:
   print "No Control of EIP\n"
 
  
 if types == "get":
  junk = generatejunk(length)
  payload = []
 
  payload.append(locateProgram(binary))
  pid = trace(payload)
  
  debug = PtraceDebugger()
  ps = debug.addProcess(pid, True)
  
  ps.cont()
  try: 
   f = open("/tmp/junk" , "w")
   f.write(junk)
   f.close()
   print "\n File /tmp/junk has been created\n"
   signal = ps.waitSignals()
   print "\n%s Detected!" % signal
   addr_eip = hex(ps.getInstrPointer())
   addr_ebp = hex(ps.getFramePointer())

   ebp = addr_ebp[8] + addr_ebp[9] + addr_ebp[6] + addr_ebp[7] + addr_ebp[4] + addr_ebp[5] + addr_ebp[2] + addr_ebp[3]


   eip = addr_eip[8] + addr_eip[9] + addr_eip[6] + addr_eip[7] + addr_eip[4] + addr_eip[5] + addr_eip[2] + addr_eip[3]
  except:
   print "\nNo SIGSEGV detected!\nMaybe the program is not vulnerable\nOr try with more bytes\n"
   os.system("rm /tmp/junk")
   debug.quit()
   sys.exit()

  os.system("rm /tmp/junk")
  try:
   print "\nControl of EBP after " + str(junk.index(ebp.decode("hex"))) + " Bytes"
  except:
   print "\nNo Control of EBP"


  try:
   print "Control of EIP after " + str(junk.index(eip.decode("hex"))) + " Bytes\n"
  except:
   print "No Control of EIP\n"

if len(sys.argv) == 4:
 checkfault(sys.argv[2], sys.argv[1], int(sys.argv[3]))
else:
 print "\nUsage: %s [Type] [Binary] [Junk_Length]\nType:\n arg = If the binary use argv\n get = If the binary use gets(), scanf() or other\n" % sys.argv[0]