poke-gate 0.1.9 → 0.2.1

package/src/permission-service.js

@@ -0,0 +1,128 @@
+import { createHash, createHmac, randomUUID } from "node:crypto";
+
+const SAFE_TOOLS = new Set(["read_file", "read_image", "list_directory", "system_info"]);
+const RISKY_TOOLS = new Set(["run_command", "write_file", "take_screenshot"]);
+
+const TOKEN_TTL_MS = 5 * 60 * 1000;
+
+function nowMs(nowFn) {
+  const value = nowFn();
+  if (value instanceof Date) return value.getTime();
+  return Number(value);
+}
+
+function normalizeForStableJson(value) {
+  if (value === null || value === undefined) return value;
+  if (Array.isArray(value)) return value.map(normalizeForStableJson);
+  if (value instanceof Date) return value.toISOString();
+  if (typeof value === "object") {
+    const normalized = {};
+    for (const key of Object.keys(value).sort()) {
+      normalized[key] = normalizeForStableJson(value[key]);
+    }
+    return normalized;
+  }
+  if (typeof value === "bigint") return value.toString();
+  return value;
+}
+
+function hashArgs(toolArgs) {
+  const stable = JSON.stringify(normalizeForStableJson(toolArgs ?? {}));
+  return createHash("sha256").update(stable).digest("hex");
+}
+
+function escapeRegex(text) {
+  return text.replace(/[.*+?^${}()|[\]\\]/g, "\\$&");
+}
+
+function patternMatches(pattern, commandText) {
+  const regexSource = `^${pattern.split("*").map(escapeRegex).join(".*")}$`;
+  return new RegExp(regexSource).test(commandText);
+}
+
+export class PermissionService {
+  constructor({ secret, now } = {}) {
+    if (!secret) throw new Error("PermissionService requires a secret");
+
+    this.secret = secret;
+    this.now = typeof now === "function" ? now : () => Date.now();
+    this.pendingApprovals = new Map();
+    this.sessionWhitelist = new Map();
+  }
+
+  isRisky(toolName) {
+    return RISKY_TOOLS.has(toolName);
+  }
+
+  requestApproval(sessionId, toolName, toolArgs) {
+    const argsHash = hashArgs(toolArgs);
+    const approvalRequestId = randomUUID();
+    const expiresAt = nowMs(this.now) + TOKEN_TTL_MS;
+    const tokenPayload = `${approvalRequestId}:${sessionId}:${toolName}:${argsHash}:${expiresAt}`;
+    const token = createHmac("sha256", this.secret).update(tokenPayload).digest("hex");
+
+    this.pendingApprovals.set(token, {
+      approvalRequestId,
+      sessionId,
+      toolName,
+      argsHash,
+      expiresAt,
+      consumed: false,
+    });
+
+    return { approvalRequestId, token, expiresAt };
+  }
+
+  validateApprovalToken(sessionId, token, toolName, toolArgs) {
+    const record = this.pendingApprovals.get(token);
+    if (!record) return false;
+    if (record.consumed) return false;
+
+    if (nowMs(this.now) > record.expiresAt) {
+      this.pendingApprovals.delete(token);
+      return false;
+    }
+
+    const argsHash = hashArgs(toolArgs);
+    if (
+      record.sessionId !== sessionId ||
+      record.toolName !== toolName ||
+      record.argsHash !== argsHash
+    ) {
+      return false;
+    }
+
+    record.consumed = true;
+    return true;
+  }
+
+  allowPatternForSession(sessionId, pattern) {
+    if (!this.sessionWhitelist.has(sessionId)) {
+      this.sessionWhitelist.set(sessionId, new Set());
+    }
+    this.sessionWhitelist.get(sessionId).add(pattern);
+  }
+
+  isAllowedBySessionPattern(sessionId, commandText) {
+    const patterns = this.sessionWhitelist.get(sessionId);
+    if (!patterns) return false;
+
+    for (const pattern of patterns) {
+      if (patternMatches(pattern, commandText)) return true;
+    }
+
+    return false;
+  }
+
+  clearSession(sessionId) {
+    this.sessionWhitelist.delete(sessionId);
+
+    for (const [token, record] of this.pendingApprovals.entries()) {
+      if (record.sessionId === sessionId) {
+        this.pendingApprovals.delete(token);
+      }
+    }
+  }
+}
+
+export { SAFE_TOOLS, RISKY_TOOLS, TOKEN_TTL_MS };

package/test/mcp-server-access-policy.test.js

@@ -0,0 +1,40 @@
+import assert from "node:assert/strict";
+import test from "node:test";
+
+import { evaluateAccessPolicy } from "../src/mcp-server.js";
+
+test("full mode allows all current tools", () => {
+  assert.equal(evaluateAccessPolicy("write_file", { path: "~/a.txt", content: "x" }, "full"), null);
+  assert.equal(evaluateAccessPolicy("take_screenshot", {}, "full"), null);
+  assert.equal(evaluateAccessPolicy("run_command", { command: "sudo reboot" }, "full"), null);
+});
+
+test("limited mode allows curated operational commands", () => {
+  assert.equal(evaluateAccessPolicy("run_command", { command: "yt-dlp https://youtu.be/demo" }, "limited"), null);
+  assert.equal(evaluateAccessPolicy("run_command", { command: "curl -I https://example.com" }, "limited"), null);
+  assert.equal(evaluateAccessPolicy("network_speed", {}, "limited"), null);
+});
+
+test("limited mode blocks dangerous or restricted actions", () => {
+  const dangerous = evaluateAccessPolicy("run_command", { command: "sudo rm -rf /" }, "limited");
+  assert.equal(typeof dangerous, "string");
+
+  const ffmpegDenied = evaluateAccessPolicy("run_command", { command: "ffmpeg -i in.mkv out.mp4" }, "limited");
+  assert.equal(typeof ffmpegDenied, "string");
+
+  const ddDenied = evaluateAccessPolicy("run_command", { command: "dd if=/dev/zero of=/tmp/a bs=1m count=1" }, "limited");
+  assert.equal(typeof ddDenied, "string");
+
+  const restrictedTool = evaluateAccessPolicy("write_file", { path: "~/a.txt", content: "x" }, "limited");
+  assert.equal(restrictedTool, "This tool is disabled in Limited Permissions mode.");
+});
+
+test("sandbox mode allows broader command set under OS sandbox", () => {
+  assert.equal(evaluateAccessPolicy("run_command", { command: "ffprobe -version" }, "sandbox"), null);
+  assert.equal(evaluateAccessPolicy("run_command", { command: "brew install yt-dlp" }, "sandbox"), null);
+  assert.equal(evaluateAccessPolicy("run_command", { command: "dd if=/dev/zero of=/tmp/a bs=1m count=1" }, "sandbox"), null);
+  assert.equal(evaluateAccessPolicy("network_speed", {}, "sandbox"), null);
+
+  const screenshotDenied = evaluateAccessPolicy("take_screenshot", {}, "sandbox");
+  assert.equal(screenshotDenied, "This tool is disabled in Sandbox mode.");
+});

package/test/mcp-server-loop-guard.test.js

@@ -0,0 +1,57 @@
+import assert from "node:assert/strict";
+import test from "node:test";
+
+import {
+  prepareRunCommandAttempt,
+  recordRunCommandOutcome,
+  resetRunCommandLoopGuard,
+} from "../src/mcp-server.js";
+
+const sessionId = "session-1";
+const cleanArgs = { command: "while ps -p 58675 58676 > /dev/null 2>&1; do sleep 5; done", cwd: "~" };
+
+test("run_command guard suppresses in-flight duplicates", () => {
+  resetRunCommandLoopGuard();
+
+  const first = prepareRunCommandAttempt(sessionId, cleanArgs, 1_000);
+  const second = prepareRunCommandAttempt(sessionId, cleanArgs, 1_001);
+
+  assert.equal(first.suppressed, false);
+  assert.equal(second.suppressed, true);
+  assert.equal(second.reason, "already_running");
+});
+
+test("run_command guard suppresses repeats after failure and clears after success", () => {
+  resetRunCommandLoopGuard();
+
+  const initial = prepareRunCommandAttempt(sessionId, cleanArgs, 2_000);
+  assert.equal(initial.suppressed, false);
+
+  recordRunCommandOutcome(sessionId, cleanArgs, { exitCode: 1 }, 2_010);
+
+  const afterFailure = prepareRunCommandAttempt(sessionId, cleanArgs, 2_020);
+  assert.equal(afterFailure.suppressed, true);
+  assert.equal(afterFailure.reason, "recent_failure");
+
+  const afterCooldown = prepareRunCommandAttempt(sessionId, cleanArgs, 2_010 + 60_000 + 1);
+  assert.equal(afterCooldown.suppressed, false);
+
+  recordRunCommandOutcome(sessionId, cleanArgs, { exitCode: 0 }, 62_020);
+  const afterSuccess = prepareRunCommandAttempt(sessionId, cleanArgs, 62_021);
+  assert.equal(afterSuccess.suppressed, false);
+});
+
+test("run_command guard scopes suppression to the same command and cwd", () => {
+  resetRunCommandLoopGuard();
+
+  const first = prepareRunCommandAttempt(sessionId, cleanArgs, 3_000);
+  assert.equal(first.suppressed, false);
+
+  recordRunCommandOutcome(sessionId, cleanArgs, { exitCode: 1 }, 3_010);
+
+  const differentCwd = prepareRunCommandAttempt(sessionId, { ...cleanArgs, cwd: "/tmp" }, 3_020);
+  assert.equal(differentCwd.suppressed, false);
+
+  const differentSession = prepareRunCommandAttempt("session-2", cleanArgs, 3_020);
+  assert.equal(differentSession.suppressed, false);
+});

package/test/mcp-server-sandbox-command.test.js

@@ -0,0 +1,18 @@
+import assert from "node:assert/strict";
+import test from "node:test";
+
+import { buildSandboxWrappedCommand } from "../src/mcp-server.js";
+
+test("buildSandboxWrappedCommand wraps command with sandbox-exec", () => {
+  const wrapped = buildSandboxWrappedCommand("ffmpeg -i in.mkv out.mp4");
+
+  assert.equal(wrapped.includes("/usr/bin/sandbox-exec -p"), true);
+  assert.equal(wrapped.includes("/bin/zsh -lc"), true);
+  assert.equal(wrapped.includes("ffmpeg -i in.mkv out.mp4"), true);
+});
+
+test("buildSandboxWrappedCommand escapes single quotes", () => {
+  const wrapped = buildSandboxWrappedCommand("echo 'hello'");
+
+  assert.equal(wrapped.includes("'\"'\"'"), true);
+});

package/test/permission-service.test.js

@@ -0,0 +1,97 @@
+import assert from "node:assert/strict";
+import test from "node:test";
+
+import { PermissionService } from "../src/permission-service.js";
+
+function createClock(startMs = 0) {
+  let current = startMs;
+  return {
+    now: () => current,
+    advance: (ms) => {
+      current += ms;
+    },
+  };
+}
+
+test("safe and risky tool classification", () => {
+  const service = new PermissionService({ secret: "test-secret" });
+
+  assert.equal(service.isRisky("read_file"), false);
+  assert.equal(service.isRisky("read_image"), false);
+  assert.equal(service.isRisky("list_directory"), false);
+  assert.equal(service.isRisky("system_info"), false);
+
+  assert.equal(service.isRisky("run_command"), true);
+  assert.equal(service.isRisky("write_file"), true);
+  assert.equal(service.isRisky("take_screenshot"), true);
+});
+
+test("session whitelist match", () => {
+  const service = new PermissionService({ secret: "test-secret" });
+
+  service.allowPatternForSession("s1", "rm /tmp/*");
+
+  assert.equal(service.isAllowedBySessionPattern("s1", "rm /tmp/file-a"), true);
+});
+
+test("whitelist isolation across sessions", () => {
+  const service = new PermissionService({ secret: "test-secret" });
+
+  service.allowPatternForSession("s1", "rm /tmp/*");
+
+  assert.equal(service.isAllowedBySessionPattern("s2", "rm /tmp/file-a"), false);
+});
+
+test("expired token invalid", () => {
+  const clock = createClock(10_000);
+  const service = new PermissionService({ secret: "test-secret", now: clock.now });
+
+  const approval = service.requestApproval("s1", "run_command", { command: "ls" });
+  clock.advance(5 * 60 * 1000 + 1);
+
+  const ok = service.validateApprovalToken("s1", approval.token, "run_command", { command: "ls" });
+  assert.equal(ok, false);
+});
+
+test("hash mismatch invalid", () => {
+  const service = new PermissionService({ secret: "test-secret" });
+
+  const approval = service.requestApproval("s1", "run_command", { command: "echo hi" });
+  const ok = service.validateApprovalToken("s1", approval.token, "run_command", { command: "echo bye" });
+
+  assert.equal(ok, false);
+});
+
+test("consumed token cannot be reused", () => {
+  const service = new PermissionService({ secret: "test-secret" });
+
+  const approval = service.requestApproval("s1", "run_command", { command: "pwd" });
+
+  const first = service.validateApprovalToken("s1", approval.token, "run_command", { command: "pwd" });
+  const second = service.validateApprovalToken("s1", approval.token, "run_command", { command: "pwd" });
+
+  assert.equal(first, true);
+  assert.equal(second, false);
+});
+
+test("glob semantics with anchored wildcard", () => {
+  const service = new PermissionService({ secret: "test-secret" });
+
+  service.allowPatternForSession("s1", "rm /tmp/*");
+
+  assert.equal(service.isAllowedBySessionPattern("s1", "rm /tmp/a"), true);
+  assert.equal(service.isAllowedBySessionPattern("s1", "rm ~/a"), false);
+});
+
+test("concurrent approvals generate unique approvalRequestId", async () => {
+  const service = new PermissionService({ secret: "test-secret" });
+
+  const approvals = await Promise.all(
+    Array.from({ length: 100 }, (_, i) =>
+      Promise.resolve(service.requestApproval("s1", "run_command", { command: `echo ${i}` })),
+    ),
+  );
+
+  const ids = new Set(approvals.map((item) => item.approvalRequestId));
+  assert.equal(ids.size, approvals.length);
+});