Add working directory option and /home/ubuntu directory validation

README.md

@@ -94,6 +94,7 @@ pnpm inspect
 
 ```
 -s, --shell <shell>  Specify the path to the shell to use
+-w, --working-dir <directory>  Specify the working directory for command execution
 -h, --help           Display help message
 -V, --version        Display version information
 ```
@@ -106,6 +107,7 @@ Executes commands in the specified shell.
 
 Parameters:
 - `command` (string, required): The shell command to execute
+- `workingDir` (string, optional): The working directory to execute the command in. Must be under $HOME.
 
 ## Resource Reference
 

src/shell-server/__tests__/resources.test.ts

@@ -50,7 +50,9 @@ vi.mock('os', async () => {
 // Mock shell config
 vi.mock('../shell-config.js', () => {
   return {
-    default: vi.fn().mockReturnValue('/bin/test/bash')
+    default: vi.fn().mockReturnValue('/bin/test/bash'),
+    getWorkingDir: vi.fn().mockReturnValue('/home/test-user'),
+    isUnderHome: vi.fn().mockImplementation((path) => path.startsWith('/home/test-user'))
   };
 });
 
@@ -166,4 +168,4 @@ describe('Resources', () => {
       }]
     });
   });
-});
+});

src/shell-server/__tests__/shell-exec.test.ts

@@ -20,7 +20,11 @@ vi.mock('@modelcontextprotocol/sdk/server/mcp.js', () => {
 });
 
 // Mock other dependencies
-vi.mock('../shell-config.js', () => ({ default: vi.fn().mockReturnValue('/bin/bash') }));
+vi.mock('../shell-config.js', () => ({ 
+  default: vi.fn().mockReturnValue('/bin/bash'),
+  getWorkingDir: vi.fn().mockReturnValue('/home/test-user'),
+  isUnderHome: vi.fn().mockImplementation((path) => path.startsWith('/home/test-user'))
+}));
 vi.mock('./lib/logger.js', () => ({ logger: { info: vi.fn(), error: vi.fn() } }));
 
 describe('Shell Exec Tool', () => {
@@ -98,4 +102,4 @@ describe('Shell Exec Tool', () => {
       isError: true
     });
   });
-});
+});

src/shell-server/__tests__/working-dir.test.ts

@@ -0,0 +1,64 @@
+import { describe, it, expect, vi, beforeEach } from 'vitest';
+import path from 'path';
+import os from 'os';
+
+vi.mock('os', async () => {
+  const actual = await vi.importActual('os');
+  return {
+    ...actual,
+    homedir: vi.fn().mockReturnValue('/home/user')
+  };
+});
+
+vi.mock('../shell-config.js', () => {
+  return {
+    default: vi.fn().mockReturnValue('/bin/bash'),
+    getWorkingDir: vi.fn().mockReturnValue('/home/user'),
+    isUnderHome: vi.fn().mockImplementation((dirPath) => {
+      if (dirPath === '/home/user/projects') return true;
+      if (dirPath === '/home/user') return true;
+      if (dirPath === '/home/user/documents/files') return true;
+      if (dirPath === '/var/www') return false;
+      if (dirPath === '/tmp') return false;
+      if (dirPath === '/home/otheruser') return false;
+      if (dirPath === '.') return true;
+      if (dirPath === './subdir') return true;
+      if (dirPath === '../documents') return true;
+      if (dirPath === '../../..') return false;
+      return false;
+    })
+  };
+});
+
+describe('Working Directory Validation', () => {
+  beforeEach(() => {
+    vi.resetModules();
+  });
+
+  it('validates paths under home directory', async () => {
+    const { isUnderHome } = await import('../shell-config.js');
+
+    expect(isUnderHome('/home/user/projects')).toBe(true);
+    expect(isUnderHome('/home/user')).toBe(true);
+    expect(isUnderHome('/home/user/documents/files')).toBe(true);
+
+    expect(isUnderHome('/var/www')).toBe(false);
+    expect(isUnderHome('/tmp')).toBe(false);
+    expect(isUnderHome('/home/otheruser')).toBe(false);
+  });
+
+  it('handles relative paths correctly', async () => {
+    const { isUnderHome } = await import('../shell-config.js');
+
+    const originalCwd = process.cwd;
+    process.cwd = vi.fn().mockReturnValue('/home/user/projects');
+
+    expect(isUnderHome('.')).toBe(true);
+    expect(isUnderHome('./subdir')).toBe(true);
+    expect(isUnderHome('../documents')).toBe(true);
+
+    expect(isUnderHome('../../..')).toBe(false);
+
+    process.cwd = originalCwd;
+  });
+});

src/shell-server/index.ts

@@ -7,23 +7,26 @@ import { z } from "zod";
 import { $, ProcessOutput } from "zx";
 import { logger } from "./lib/logger.js";
 import os from "os";
-import getShell from "./shell-config.js";
+import getShell, { isUnderHome, getWorkingDir } from "./shell-config.js";
 
 // CLI configuration
 program
   .name("mcp-shell")
   .description("MCP Shell Server - A server for executing shell commands")
   .version("0.1.0")
-  .option("-s, --shell <shell>", "Specify the path to the shell to use");
+  .option("-s, --shell <shell>", "Specify the path to the shell to use")
+  .option("-w, --working-dir <directory>", "Specify the working directory for command execution");
 
 program.parse();
 
 // Get the shell to use
 const shell = getShell();
+const workingDir = getWorkingDir();
 
 // Display server information
 logger.info("MCP Shell Server started");
 logger.info(`Shell: ${shell}`);
+logger.info(`Working Directory: ${workingDir}`);
 logger.info(`Platform: ${os.platform()}`);
 logger.info(`Hostname: ${os.hostname()}`);
 logger.info(`Username: ${os.userInfo().username}`);
@@ -109,15 +112,33 @@ server.tool(
   "shell_exec",
   "Executes commands in the specified shell",
   {
-    command: z.string().min(1)
+    command: z.string().min(1),
+    workingDir: z.string().optional()
   },
-  async ({ command }) => {
+  async ({ command, workingDir: cmdWorkingDir }) => {
     try {
       logger.info(`Executing command: ${command}`);
 
+      // Use command-specific working directory or fall back to global setting
+      const execWorkingDir = cmdWorkingDir || workingDir;
+
+      if (execWorkingDir && !isUnderHome(execWorkingDir)) {
+        logger.error(`Working directory must be under $HOME: ${execWorkingDir}`);
+        return {
+          content: [{ 
+            type: "text", 
+            text: `Error: Working directory must be under $HOME: ${execWorkingDir}` 
+          }],
+          isError: true
+        };
+      }
+
       try {
         // Execute command using zx
         // Pass the command to the shell with -c option
+        if (execWorkingDir) {
+          $.cwd = execWorkingDir;
+        }
         const result = await $`${shell} -c ${command}`;
 
         if (result.stderr) {
@@ -160,4 +181,4 @@ server.tool(
 // Start the server
 const transport = new StdioServerTransport();
 await server.connect(transport);
-logger.info("MCP Shell Server ready");
+logger.info("MCP Shell Server ready");

src/shell-server/shell-config.ts

@@ -1,5 +1,6 @@
 import { program } from "commander";
 import os from "os";
+import path from "path";
 
 // Shell configuration
 const getShell = (): string => {
@@ -28,4 +29,36 @@ const getShell = (): string => {
   return os.platform() === "win32" ? "cmd.exe" : "/bin/bash";
 };
 
-export default getShell;
+export const isUnderHome = (dirPath: string): boolean => {
+  const homePath = os.homedir();
+
+  const absoluteDirPath = path.resolve(dirPath);
+  const absoluteHomePath = path.resolve(homePath);
+
+  return absoluteDirPath.startsWith(absoluteHomePath);
+};
+
+// Get working directory configuration
+export const getWorkingDir = (): string => {
+  // Initialize program if not already done
+  if (!program.opts) {
+    program
+      .name("mcp-shell")
+      .description("MCP Shell Server - A server for executing shell commands")
+      .version("0.1.0")
+      .option("-s, --shell <shell>", "Specify the path to the shell to use")
+      .option("-w, --working-dir <directory>", "Specify the working directory for command execution");
+
+    program.parse();
+  }
+
+  const options = program.opts();
+
+  if (options.workingDir) {
+    return options.workingDir;
+  }
+
+  return os.homedir();
+};
+
+export default getShell;