terminal

import { tool, Tool, ToolsProviderController } from "@lmstudio/sdk";
import { z } from "zod";
import { execSync } from "child_process";
import * as fs from "fs";
import * as path from "path";
import * as os from "os";
import * as crypto from "crypto";
import * as net from "net";

const MAX_OUTPUT_LENGTH = 100000;

function truncateOutput(output: string): string {
  if (output.length > MAX_OUTPUT_LENGTH) {
    return output.slice(0, MAX_OUTPUT_LENGTH) + `\n\n[Output truncated - ${output.length} chars total]`;
  }
  return output;
}

function sanitizeCommand(command: string): { valid: boolean; command?: string; error?: string } {
  // Block dangerous shell metacharacters that could escape the sandbox
  const dangerousPatterns = [
    /\b(rm\s+-rf|mkfs|fdisk|dd\s+if=|shred|wipe)\b/i,
    /;\s*(rm|sudo|su\s+root|chmod\s+777|chown\s+root)\b/i,
    /\|\s*(rm|sudo|su\s+root|chmod\s+777|chown\s+root)\b/i,
    /`[^`]*\b(rm|sudo|su\s+root)\b[^`]*`/,
    /\$\([^)]*\b(rm|sudo|su\s+root)\b[^)]*\)/,
    /&&\s*(rm|sudo|su\s+root)\b/i,
    /\|\|\s*(rm|sudo|su\s+root)\b/i,
    />\s*\/dev\/(sda|sdb|vda|vdb|hda|hdb)/i,
    /<\s*\/dev\/(sda|sdb|vda|vdb|hda|hdb)/i,
    /\/etc\/(passwd|shadow|sudoers)\b/i,
    /\/proc\/(kcore|self\/mem)\b/i,
    /\/sys\/(firmware|device)\b/i,
    new RegExp("/dev/(random|urandom)\\b(?=.*\\/>)", "i"),
    /\/dev\/(null|zero)\b(?=.*<)/i,
    /\/boot\/(grub|loader|vmlinuz|initrd)/i,
    /\/lib(64)?\/(modules|firmware)/i,
    /\/usr\/lib(64)?\/(firmware|systemd)/i,
    /\/var\/(lib\/dpkg|log\/auth)/i,
    /\/root\/\./i,
    /\/etc\/(pam\.d|ssh\/ssh_config|crontab)/i,
    /\/proc\/(kallsyms|kexec_load|modules)\b/i,
    /\/sys\/(kernel|security|fs)/i,
    /\/dev\/shm\//i,
    /\/run\/(systemd|docker)/i,
    /\/var\/run\//i,
    /\/etc\/init\.d\//i,
    /\/usr\/local\/(bin|sbin)\//i,
    /\/opt\//i,
    /\/snap\//i,
    /\/etc\/ld\.so\./i,
    /\/etc\/nsswitch\.conf/i,
    /\/etc\/hosts\.allow/i,
    /\/etc\/hosts\.deny/i,
    /\/etc\/security\//i,
    /\/etc\/apparmor\//i,
    /\/etc\/selinux\//i,
    /\/etc\/modprobe\./i,
    /\/etc\/modprobe\.d\//i,
    /\/etc\/udev\//i,
    /\/etc\/systemd\//i,
    /\/etc\/cron\./i,
    /\/etc\/logrotate\./i,
    /\/etc\/sysctl\./i,
    /\/etc\/fstab/i,
    /\/etc\/mtab/i,
    /\/etc\/group\b/i,
    /\/etc\/gshadow/i,
    /\/etc\/shadow-$/i,
    /\/etc\/passwd-$/i,
    /\/etc\/sudoers\.d\//i,
    /\/etc\/ssh\/ssh_host_/i,
    /\/etc\/ssl\/private\//i,
    /\/etc\/pki\//i,
    /\/etc\/certs\//i,
    /\/etc\/trust\//i,
    /\/etc\/ca-certificates\//i,
    /\/etc\/pam\.d\//i,
    /\/etc\/security\//i,
    /\/etc\/security\/limits\.conf/i,
    /\/etc\/security\/limits\.d\//i,
    /\/etc\/security\/pam_env\.conf/i,
    /\/etc\/security\/pam_access\.conf/i,
    /\/etc\/security\/pam_env\.conf/i,
    /\/etc\/security\/pam_limits\.conf/i,
    /\/etc\/security\/pam_namespace\.conf/i,
    /\/etc\/security\/pam_time\.conf/i,
    /\/etc\/security\/pam_unix\.conf/i,
    /\/etc\/security\/pam_winbind\.conf/i,
    /\/etc\/security\/pam_ssh_agent_auth\.conf/i,
    /\/etc\/security\/pam_ssh_keygen\.conf/i,
    /\/etc\/security\/pam_ssh_keygen\.d\//i,
    /\/etc\/security\/pam_ssh_agent_auth\.d\//i,
    /\/etc\/security\/pam_time\.d\//i,
    /\/etc\/security\/pam_limits\.d\//i,
    /\/etc\/security\/pam_namespace\.d\//i,
    /\/etc\/security\/pam_env\.d\//i,
    /\/etc\/security\/pam_access\.d\//i,
    /\/etc\/security\/pam_unix\.d\//i,
    /\/etc\/security\/pam_winbind\.d\//i,
    /\/etc\/security\/pam_ssh_agent_auth\.d\//i,
    /\/etc\/security\/pam_ssh_keygen\.d\//i,
    /\/etc\/security\/pam_time\.d\//i,
    /\/etc\/security\/pam_limits\.d\//i,
    /\/etc\/security\/pam_namespace\.d\//i,
    /\/etc\/security\/pam_env\.d\//i,
    /\/etc\/security\/pam_access\.d\//i,
    /\/etc\/security\/pam_unix\.d\//i,
    /\/etc\/security\/pam_winbind\.d\//i,
  ];

  for (const pattern of dangerousPatterns) {
    if (pattern.test(command)) {
      return { valid: false, error: `Command blocked: potentially dangerous operation detected` };
    }
  }

  // Allow safe commands
  const safeCommands = [
    /^ls\b/, /^cat\b/, /^echo\b/, /^pwd\b/, /^cd\b/, /^mkdir\b/,
    /^rm\b/, /^cp\b/, /^mv\b/, /^find\b/, /^grep\b/, /^head\b/,
    /^tail\b/, /^wc\b/, /^sort\b/, /^uniq\b/, /^diff\b/, /^sed\b/,
    /^awk\b/, /^cut\b/, /^tr\b/, /^date\b/, /^whoami\b/, /^id\b/,
    /^hostname\b/, /^uname\b/, /^df\b/, /^du\b/, /^free\b/, /^top\b/,
    /^ps\b/, /^netstat\b/, /^ss\b/, /^ping\b/, /^curl\b/, /^wget\b/,
    /^python\b/, /^python3\b/, /^node\b/, /^npm\b/, /^pip\b/, /^pip3\b/,
    /^git\b/, /^ssh\b/, /^scp\b/, /^rsync\b/, /^tar\b/, /^gzip\b/,
    /^gunzip\b/, /^zip\b/, /^unzip\b/, /^chown\b/, /^chmod\b/,
    /^chmod\b/, /^chgrp\b/, /^stat\b/, /^file\b/, /^ln\b/, /^touch\b/,
    /^which\b/, /^whereis\b/, /^man\b/, /^info\b/, /^help\b/,
    /^clear\b/, /^reset\b/, /^tput\b/, /^stty\b/, /^kill\b/,
    /^killall\b/, /^pkill\b/, /^xargs\b/, /^tee\b/, /^xargs\b/,
    /^bc\b/, /^dc\b/, /^expr\b/, /^seq\b/, /^yes\b/,
    /^base64\b/, /^md5sum\b/, /^sha256sum\b/, /^sha1sum\b/,
    /^openssl\b/, /^gpg\b/, /^gpg2\b/,
    /^apt\b/, /^apt-get\b/, /^dpkg\b/, /^yum\b/, /^dnf\b/,
    /^rpm\b/, /^pacman\b/, /^zypper\b/, /^apk\b/,
    /^brew\b/, /^choco\b/, /^winget\b/,
    /^docker\b/, /^kubectl\b/, /^helm\b/, /^terraform\b/,
    /^ansible\b/, /^kubectl\b/, /^helm\b/, /^terraform\b/,
    /^ansible\b/, /^ansible-playbook\b/,
    /^gcc\b/, /^g\+\+/, /^clang\b/, /^clang\+\+/,
    /^make\b/, /^cmake\b/, /^ninja\b/, /^meson\b/, /^cargo\b/,
    /^rustc\b/, /^go\b/, /^javac\b/, /^java\b/, /^python\b/,
    /^python3\b/, /^node\b/, /^npm\b/, /^yarn\b/, /^pnpm\b/,
    /^deno\b/, /^bun\b/, /^ruby\b/, /^irb\b/, /^rails\b/,
    /^perl\b/, /^php\b/, /^ruby\b/, /^lua\b/, /^gcc\b/,
    /^g\+\+/, /^clang\b/, /^clang\+\+/, /^rustc\b/,
    /^cargo\b/, /^go\b/, /^javac\b/, /^java\b/, /^rustc\b/,
    /^cargo\b/, /^go\b/, /^javac\b/, /^java\b/, /^rustc\b/,
    /^cargo\b/, /^go\b/, /^javac\b/, /^java\b/,
  ];

  const commandName = command.trim().split(/\s+/)[0].toLowerCase();
  let isSafe = false;
  for (const pattern of safeCommands) {
    if (pattern.test(commandName)) {
      isSafe = true;
      break;
    }
  }

  if (!isSafe) {
    return { valid: false, error: `Command not allowed: ${commandName}` };
  }

  return { valid: true, command };
}

function executeShellCommand(command: string, timeoutMs: number = 30000): { success: boolean; output?: string; error?: string } {
  try {
    const sanitized = sanitizeCommand(command);
    if (!sanitized.valid) {
      return { success: false, error: sanitized.error || "Invalid command" };
    }

    const result = execSync(sanitized.command!, {
      timeout: timeoutMs,
      encoding: "utf-8",
      maxBuffer: 1024 * 1024 * 10,
      stdio: ["pipe", "pipe", "pipe"],
    });

    return { success: true, output: truncateOutput(result) };
  } catch (err: any) {
    return { success: false, error: `Shell execution failed: ${err.message || "Unknown error"}` };
  }
}

function readSpecificFile(filePath: string): { success: boolean; content?: string; output?: string; error?: string } {
  try {
    const resolved = path.resolve(filePath);
    if (!fs.existsSync(resolved)) {
      return { success: false, error: `File not found: ${filePath}` };
    }

    const stat = fs.statSync(resolved);
    if (stat.isDirectory()) {
      return { success: false, error: `Path is a directory: ${filePath}` };
    }

    const content = fs.readFileSync(resolved, "utf-8");
    return { success: true, content, output: `File read successfully (${content.length} bytes)` };
  } catch (err: any) {
    return { success: false, error: `Read failed: ${err.message}` };
  }
}

function writeSpecificFile(filePath: string, content: string): { success: boolean; output?: string; error?: string } {
  try {
    const resolvedPath = path.resolve(filePath);
    const dir = path.dirname(resolvedPath);

    if (!fs.existsSync(dir)) {
      fs.mkdirSync(dir, { recursive: true });
    }

    fs.writeFileSync(resolvedPath, content, "utf-8");
    return { success: true, output: `File written successfully (${content.length} bytes)` };
  } catch (err: any) {
    return { success: false, error: `Write failed: ${err.message}` };
  }
}

function diffFiles(file1: string, file2: string): { success: boolean; output?: string; error?: string } {
  try {
    const resolved1 = path.resolve(file1);
    const resolved2 = path.resolve(file2);

    if (!fs.existsSync(resolved1)) return { success: false, error: `File not found: ${file1}` };
    if (!fs.existsSync(resolved2)) return { success: false, error: `File not found: ${file2}` };

    const content1 = fs.readFileSync(resolved1, "utf-8");
    const content2 = fs.readFileSync(resolved2, "utf-8");

    const lines1 = content1.split("\n");
    const lines2 = content2.split("\n");

    const diffLines: string[] = [];
    const maxLines = Math.max(lines1.length, lines2.length);
    let hasChanges = false;

    for (let i = 0; i < maxLines; i++) {
      const l1 = i < lines1.length ? lines1[i] : undefined;
      const l2 = i < lines2.length ? lines2[i] : undefined;

      if (l1 === l2) {
        if (i < 3 || i >= maxLines - 3 || (i > 0 && lines1[i - 1] !== lines2[i - 1] && i < maxLines - 3)) {
          diffLines.push(` ${l1}`);
        }
      } else {
        hasChanges = true;
        diffLines.push(l1 !== undefined ? `-${l1}` : `< ${l2}`);
        diffLines.push(l2 !== undefined ? `+${l2}` : `> ${l1}`);
      }
    }

    if (!hasChanges) return { success: true, output: "Files are identical" };

    return { success: true, output: truncateOutput(diffLines.join("\n")) };
  } catch (err: any) {
    return { success: false, error: `Diff failed: ${err.message}` };
  }
}

function calculateChecksum(filePath: string, algorithm: "sha256" | "md5" = "sha256"): Promise<{ success: boolean; output?: string; error?: string }> {
  return new Promise((resolve) => {
    try {
      const resolved = path.resolve(filePath);
      if (!fs.existsSync(resolved)) {
        return resolve({ success: false, error: `File not found: ${filePath}` });
      }

      const hash = crypto.createHash(algorithm);
      const stream = fs.createReadStream(resolved);

      stream.on("data", (chunk: Buffer) => hash.update(chunk));
      stream.on("end", () => resolve({ success: true, output: hash.digest("hex") }));
      stream.on("error", (err: NodeJS.ErrnoException) => resolve({ success: false, error: `Checksum failed: ${err.message}` }));
    } catch (err: any) {
      resolve({ success: false, error: `Checksum failed: ${err.message}` });
    }
  });
}

function getSystemInformation(): { success: boolean; output?: string; error?: string } {
  try {
    const totalMem = os.totalmem();
    const freeMem = os.freemem();
    const usedPct = ((totalMem - freeMem) / totalMem * 100).toFixed(1);
    
    const result = [
      `Platform: ${process.platform}`,
      `Architecture: ${process.arch}`,
      `CPU Count: ${os.cpus().length}`,
      `Total Memory: ${(totalMem / (1024 ** 3)).toFixed(2)} GB`,
      `Free Memory: ${(freeMem / (1024 ** 3)).toFixed(2)} GB`,
      `Used Memory: ${usedPct}%`,
      `Uptime: ${os.uptime().toFixed(0)} seconds`,
      `Hostname: ${os.hostname()}`,
      `Temp Dir: ${os.tmpdir()}`,
      `Home Dir: ${os.homedir()}`,
    ].join("\n");

    return { success: true, output: result };
  } catch (err: any) {
    return { success: false, error: `System info failed: ${err.message}` };
  }
}

function listAllProcesses(): { success: boolean; output?: string; error?: string } {
  try {
    const psOutput = execSync("tasklist /FI \"MEMGT 0\" /FO CSV /NH", { encoding: "utf-8" });
    const lines = psOutput.trim().split("\n");
    
    const result = lines
      .slice(0, 50)
      .map((line) => line.replace(/"/g, "").split(","))
      .filter((parts) => parts.length >= 3)
      .map((parts) => `[PID: ${parts[1]}] ${parts[0]} - MEM: ${parts[3]}`)
      .join("\n");

    return { success: true, output: result + (lines.length > 50 ? `\n\n[... and ${lines.length - 50} more processes]` : "") };
  } catch (err: any) {
    return { success: false, error: `Process listing failed: ${err.message}` };
  }
}

function checkNetworkConnectivity(host: string, port: number = 80): Promise<{ success: boolean; output?: string; error?: string }> {
  return new Promise((resolve) => {
    const start = Date.now();
    const socket = new net.Socket();

    socket.setTimeout(5000);
    
    socket.on("connect", () => {
      const latency = Date.now() - start;
      socket.destroy();
      resolve({ success: true, output: `Connected to ${host}:${port} (${latency}ms)` });
    });

    socket.on("timeout", () => {
      socket.destroy();
      resolve({ success: false, error: `Connection to ${host}:${port} timed out` });
    });

    socket.on("error", (err: Error) => {
      resolve({ success: false, error: `Connection to ${host}:${port} failed: ${err.message}` });
    });

    socket.connect(port, host);
  });
}

function getRuntimeExecutable(runtime: string): string {
  const runtimeMap: Record<string, string> = {
    "python": "python",
    "node": "node",
    "bash": "bash",
    "sh": "sh",
    "pwsh": "pwsh",
    "ruby": "ruby",
    "perl": "perl",
    "lua": "lua",
  };
  return runtimeMap[runtime] || runtime;
}

export async function toolsProvider(ctl: ToolsProviderController): Promise<Tool[]> {
  const tools: Tool[] = [];

  // shell_exec tool
  tools.push(tool({
    name: "shell_exec",
    description: "Execute a shell command and return the output",
    parameters: {
      command: z.string().describe("The shell command to execute"),
      timeout_ms: z.number().optional().describe("Timeout in milliseconds (default: 30000)"),
    },
    implementation: async ({ command, timeout_ms }) => {
      const timeout = timeout_ms ?? 30000;
      return executeShellCommand(command, timeout);
    },
  }));

  // script_run tool
  tools.push(tool({
    name: "script_run",
    description: "Run a script in a specified runtime environment",
    parameters: {
      content: z.string().describe("The script content to execute"),
      runtime: z.enum(["bash", "sh", "python", "node", "pwsh", "ruby", "perl", "lua"]).optional().describe("Runtime to use"),
    },
    implementation: async ({ content, runtime }) => {
      const scriptContent = content || "";
      const rt = runtime || "bash";
      
      if (!scriptContent) {
        return { success: false, error: "No script content provided" };
      }

      try {
        const sanitized = sanitizeCommand(scriptContent);
        if (!sanitized.valid) {
          return { success: false, error: sanitized.error || "Invalid script" };
        }

        const tempDir = os.tmpdir();
        const ext = rt === "python" ? "py" : rt === "node" ? "js" : rt === "bash" || rt === "sh" ? "sh" : "txt";
        const tempFile = path.join(tempDir, `script_${Date.now()}.${ext}`);
        fs.writeFileSync(tempFile, scriptContent, "utf-8");

        try {
          const result = execSync(`"${getRuntimeExecutable(rt)}" "${tempFile}"`, {
            timeout: 30000,
            encoding: "utf-8",
            maxBuffer: 1024 * 1024 * 10,
          });

          return { success: true, output: truncateOutput(result) };
        } finally {
          try { fs.unlinkSync(tempFile); } catch { /* ignore */ }
        }
      } catch (err: any) {
        return { success: false, error: err.message };
      }
    },
  }));

  // file_write tool
  tools.push(tool({
    name: "file_write",
    description: "Write content to a file, creating directories as needed",
    parameters: {
      path: z.string().describe("File path to write to"),
      content: z.string().describe("Content to write"),
    },
    implementation: async ({ path: filePath, content }) => {
      return writeSpecificFile(filePath, content);
    },
  }));

  // read_file tool
  tools.push(tool({
    name: "read_file",
    description: "Read the contents of a file",
    parameters: {
      path: z.string().describe("File path to read"),
    },
    implementation: async ({ path: filePath }) => {
      return readSpecificFile(filePath);
    },
  }));

  // file_diff tool
  tools.push(tool({
    name: "file_diff",
    description: "Compare two files and show the differences",
    parameters: {
      file1: z.string().describe("First file path"),
      file2: z.string().describe("Second file path"),
    },
    implementation: async ({ file1, file2 }) => {
      return diffFiles(file1, file2);
    },
  }));

  // file_checksum tool
  tools.push(tool({
    name: "file_checksum",
    description: "Calculate the checksum (hash) of a file",
    parameters: {
      path: z.string().describe("File path"),
      algorithm: z.enum(["sha256", "md5"]).optional().describe("Hash algorithm"),
    },
    implementation: async ({ path: filePath, algorithm }) => {
      return calculateChecksum(filePath, algorithm ?? "sha256");
    },
  }));

  // system_info tool
  tools.push(tool({
    name: "system_info",
    description: "Get system information (OS, memory, CPU, etc.)",
    parameters: {},
    implementation: async () => {
      return getSystemInformation();
    },
  }));

  // list_processes tool
  tools.push(tool({
    name: "list_processes",
    description: "List running processes with resource usage",
    parameters: {},
    implementation: async () => {
      return listAllProcesses();
    },
  }));

  // check_network tool
  tools.push(tool({
    name: "check_network",
    description: "Test network connectivity to a host and port",
    parameters: {
      host: z.string().describe("Host to check"),
      port: z.number().optional().describe("Port to check (default: 80)"),
    },
    implementation: async ({ host, port }) => {
      return checkNetworkConnectivity(host, port ?? 80);
    },
  }));

  return tools;
}