pentest_mcp.py

from mcp.server.fastmcp import FastMCP, Context
import logging

from tools.nmap import run_nmap
from tools.gobuster import run_gobuster

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger("security-tools-mcp")

# Create MCP server
mcp = FastMCP("PentestMCP")

# Register tools
mcp.tool()(run_nmap)
mcp.tool()(run_gobuster)

def validate_input(input_str: str) -> bool:
    """Basic validation to help prevent command injection"""
    # Disallow dangerous characters
    dangerous_chars = [';', '&', '|', '>', '<', '`', '$', '(', ')', '{', '}', '\\']
    return not any(char in input_str for char in dangerous_chars)

def validate_options(options: str) -> bool:
    """Validate command-line options"""
    # Disallow dangerous characters
    dangerous_chars = [';', '&', '|', '>', '<', '`', '$', '(', ')', '{', '}', '\\']
    return not any(char in options for char in dangerous_chars)

def validate_wordlist(wordlist: str) -> bool:
    """Validate wordlist path"""
    # Only allow wordlists in specific directories
    allowed_prefixes = ["/usr/share/wordlists/", "/opt/wordlists/"]
    dangerous_chars = [';', '&', '|', '>', '<', '`', '$', '(', ')', '{', '}', '\\']
    
    return (
        any(wordlist.startswith(prefix) for prefix in allowed_prefixes) and
        not any(char in wordlist for char in dangerous_chars)
    )

if __name__ == "__main__":
    logger.info("Starting Pentest MCP Server")
    mcp.run(transport='stdio')