connectai/src/security.ts

import * as vscode from 'vscode';
import * as path from 'path';
import * as fs from 'fs';

/**
 * Validates that a path is strictly within the workspace.
 * Prevents Path Traversal attacks by resolving real paths and checking boundaries.
 */
/**
 * Additional trusted root paths beyond the workspace.
 * Populated once from VS Code workspace folders on first call.
 */
let _trustedRoots: string[] | null = null;

function getTrustedRoots(workspaceRoot: string): string[] {
    if (_trustedRoots) return _trustedRoots;
    const roots = [path.normalize(workspaceRoot).toLowerCase()];
    // Include all open workspace folders as trusted roots
    const folders = vscode.workspace.workspaceFolders;
    if (folders) {
        for (const f of folders) {
            roots.push(path.normalize(f.uri.fsPath).toLowerCase());
        }
    }
    // Also trust the immediate parent of each root, so sibling projects under a
    // shared parent (e.g. E:\Wiki\connectai + E:\Wiki\Datacollect) are reachable
    // for read/list. Guard: never widen to a drive/filesystem root.
    for (const r of [...roots]) {
        const parent = path.normalize(path.dirname(r)).toLowerCase();
        if (parent && parent !== r && path.dirname(parent) !== parent) {
            roots.push(parent);
        }
    }
    _trustedRoots = [...new Set(roots)];
    return _trustedRoots;
}

/** Reset cached roots (useful when workspace folders change). */
export function resetTrustedRoots(): void {
    _trustedRoots = null;
}

export function validatePath(workspaceRoot: string, targetPath: string): string {
    if (!workspaceRoot) {
        throw new Error("Security Violation: Workspace root not defined.");
    }

    const absolutePath = path.resolve(workspaceRoot, targetPath);
    const normalizedTarget = path.normalize(absolutePath).toLowerCase();
    const trusted = getTrustedRoots(workspaceRoot);

    const isTrusted = trusted.some(root => normalizedTarget.startsWith(root));
    if (!isTrusted) {
        throw new Error(`Security Violation: Path traversal detected! Attempted to access ${absolutePath} which is outside allowed boundaries.`);
    }

    return absolutePath;
}

/**
 * Splits a command on top-level `&&`, ignoring `&&` that appears inside single-
 * or double-quoted strings (e.g. a commit message). Returns trimmed, non-empty parts.
 */
function splitTopLevelAnd(command: string): string[] {
    const parts: string[] = [];
    let buf = '';
    let quote: string | null = null;
    for (let i = 0; i < command.length; i++) {
        const c = command[i];
        if (quote) {
            buf += c;
            if (c === quote) { quote = null; }
            continue;
        }
        if (c === "'" || c === '"') { quote = c; buf += c; continue; }
        if (c === '&' && command[i + 1] === '&') {
            parts.push(buf);
            buf = '';
            i++; // skip the second '&'
            continue;
        }
        buf += c;
    }
    parts.push(buf);
    return parts.map(p => p.trim()).filter(p => p.length > 0);
}

/**
 * Windows PowerShell 5.1 — the default VS Code integrated terminal on Windows —
 * does not support the `&&` chaining operator (it is a hard parser error, so the
 * WHOLE command fails to run). Local models emit `&&` constantly because every
 * git/npm tutorial uses it, and a system-prompt rule alone does not reliably
 * stop a small model. So rewrite `A && B && C` into a PowerShell-native
 * conditional chain that preserves short-circuit semantics:
 *
 *   A && B && C  ->  A; if ($?) { B; if ($?) { C } }
 *
 * `$?` reflects the success of the previous command, so a failed step still
 * short-circuits the rest — important so e.g. a failed `cd` never lets `git`
 * run in the wrong directory.
 *
 * On macOS/Linux this rewrite is actively harmful — `&&` is the native
 * POSIX-shell operator and the PowerShell `if ($?) { ... }` shape is a zsh/bash
 * syntax error. So skip the rewrite entirely off-Windows.
 */
function rewriteForPowerShell(command: string): string {
    if (process.platform !== 'win32') { return command; }
    if (!command.includes('&&')) { return command; }
    const parts = splitTopLevelAnd(command);
    if (parts.length <= 1) { return command; }
    let chain = parts[parts.length - 1];
    for (let i = parts.length - 2; i >= 0; i--) {
        chain = `${parts[i]}; if ($?) { ${chain} }`;
    }
    return chain;
}

/**
 * Sanitizes terminal commands to prevent destructive actions.
 * Uses a combination of blocklist for dangerous patterns and recommendation for allowed tools.
 */
export function sanitizeCommand(command: string): string {
    const trimmedCmd = command.trim();
    
    // 1. Dangerous Shell Characters/Patterns (Blocklist)
    const dangerousPatterns = [
        /rm\s+-rf\s+\//,         // Root deletion
        /mkfs/,                  // Filesystem formatting
        /dd\s+if=/,               // Low-level disk writing
        />\s*\/dev\/sd/,         // Direct disk access
        /:(){:|:&};:/,           // Fork bomb
        /shutdown/,              // System shutdown
        /reboot/,                // System reboot
        /mv\s+.*\/dev\/null/     // Moving to null
    ];

    for (const pattern of dangerousPatterns) {
        if (pattern.test(trimmedCmd)) {
            throw new Error(`Security Violation: Destructive command pattern detected! Blocked: ${trimmedCmd}`);
        }
    }

    // 2. Allowlist of safe base commands (Optional but recommended)
    // For now, we allow common development tools
    const safeBaseCommands = [
        'npm', 'node', 'npx', 'git', 'python', 'python3', 'pip', 'pip3',
        'cargo', 'rustc', 'go', 'gcc', 'g++', 'make', 'ls', 'cat', 'echo',
        'mkdir', 'cp', 'mv', 'touch'
    ];
    
    const baseCmd = trimmedCmd.split(/\s+/)[0];
    if (baseCmd && !safeBaseCommands.includes(baseCmd)) {
        console.warn(`[Security] Warning: Running uncommon command '${baseCmd}'. Ensure this is intended.`);
    }

    // Rewrite `&&` chains for PowerShell (the Windows default terminal) so the
    // command actually runs instead of failing with a parser error.
    return rewriteForPowerShell(trimmedCmd);
}