picogent/src/sandbox.ts

import path from 'path';
import fs from 'fs/promises';
import {
  createReadTool,
  createWriteTool,
  createEditTool,
  createBashTool,
  createGrepTool,
  createFindTool,
  createLsTool,
} from '@mariozechner/pi-coding-agent';
import type {
  ReadOperations,
  WriteOperations,
  EditOperations,
} from '@mariozechner/pi-coding-agent';
import type { AgentTool } from '@mariozechner/pi-agent-core';

/**
 * Resolve and normalize a path for comparison.
 * Handles ~, relative paths, and normalizes separators.
 */
function resolvePath(filePath: string, cwd: string): string {
  let expanded = filePath;
  if (expanded.startsWith('~')) {
    expanded = path.join(process.env.HOME || process.env.USERPROFILE || '', expanded.slice(1));
  }
  if (!path.isAbsolute(expanded)) {
    expanded = path.resolve(cwd, expanded);
  }
  return path.normalize(expanded);
}

/**
 * Check if a resolved absolute path is within any of the allowed directories.
 */
function isPathAllowed(absolutePath: string, allowedPaths: string[]): boolean {
  const normalized = path.normalize(absolutePath) + path.sep;
  for (const allowed of allowedPaths) {
    const allowedNorm = path.normalize(allowed) + path.sep;
    if (normalized.startsWith(allowedNorm) || path.normalize(absolutePath) === path.normalize(allowed)) {
      return true;
    }
  }
  return false;
}

function pathError(absolutePath: string, allowedPaths: string[]): Error {
  return new Error(
    `Access denied: "${absolutePath}" is outside allowed paths.\n` +
    `Allowed: ${allowedPaths.join(', ')}`
  );
}

/**
 * Create a set of coding tools with path restrictions.
 *
 * - If `allowedPaths` is provided and non-empty, all file operations (read, write, edit)
 *   are validated to ensure they only access files within those directories.
 *   Bash commands are restricted to run with cwd inside allowed paths.
 *
 * - If `allowedPaths` is empty or undefined, tools work without restrictions
 *   (same as default codingTools, but bound to the given cwd).
 */
export function createSandboxedTools(cwd: string, allowedPaths?: string[]): AgentTool<any>[] {
  // Resolve allowed paths to absolute
  const resolved = (allowedPaths || [])
    .map(p => path.resolve(cwd, p))
    .filter(p => p.length > 0);

  // No restrictions — just create tools bound to cwd
  if (resolved.length === 0) {
    return [
      createReadTool(cwd),
      createBashTool(cwd),
      createEditTool(cwd),
      createWriteTool(cwd),
      createGrepTool(cwd),
      createFindTool(cwd),
      createLsTool(cwd),
    ];
  }

  // --- Sandboxed operations ---

  function guard(absolutePath: string): void {
    const full = resolvePath(absolutePath, cwd);
    if (!isPathAllowed(full, resolved)) {
      throw pathError(full, resolved);
    }
  }

  const readOps: ReadOperations = {
    readFile: async (absolutePath: string) => {
      guard(absolutePath);
      return fs.readFile(absolutePath);
    },
    access: async (absolutePath: string) => {
      guard(absolutePath);
      await fs.access(absolutePath);
    },
  };

  const writeOps: WriteOperations = {
    writeFile: async (absolutePath: string, content: string) => {
      guard(absolutePath);
      await fs.writeFile(absolutePath, content, 'utf-8');
    },
    mkdir: async (dir: string) => {
      guard(dir);
      await fs.mkdir(dir, { recursive: true });
    },
  };

  const editOps: EditOperations = {
    readFile: async (absolutePath: string) => {
      guard(absolutePath);
      return fs.readFile(absolutePath);
    },
    writeFile: async (absolutePath: string, content: string) => {
      guard(absolutePath);
      await fs.writeFile(absolutePath, content, 'utf-8');
    },
    access: async (absolutePath: string) => {
      guard(absolutePath);
      await fs.access(absolutePath);
    },
  };

  return [
    createReadTool(cwd, { operations: readOps }),
    createBashTool(cwd, {
      spawnHook: (ctx) => {
        // Ensure bash cwd is within allowed paths
        const bashCwd = resolvePath(ctx.cwd, cwd);
        if (!isPathAllowed(bashCwd, resolved)) {
          throw pathError(bashCwd, resolved);
        }
        return ctx;
      },
    }),
    createEditTool(cwd, { operations: editOps }),
    createWriteTool(cwd, { operations: writeOps }),
    createGrepTool(cwd),
    createFindTool(cwd),
    createLsTool(cwd),
  ];
}