__init__.py•2.7 kB
"""
Tools Package
Auto-discovers and loads all security tools (MCP-compatible)
"""
import os
import sys
import importlib
from typing import Dict, List
from pathlib import Path
# Add parent directory to path
sys.path.insert(0, str(Path(__file__).parent.parent))
from tools.base_tool import BaseTool
from mcp.types import MCPTool
class ToolRegistry:
"""Registry for managing security tools (MCP-compatible)"""
def __init__(self):
self.tools: Dict[str, BaseTool] = {}
async def load_all_tools(self) -> int:
"""Auto-discover and load all tools from this directory"""
tools_dir = os.path.dirname(__file__)
loaded_count = 0
for filename in os.listdir(tools_dir):
if filename.endswith('.py') and filename not in ['__init__.py', 'base_tool.py']:
tool_name = filename[:-3] # Remove .py extension
try:
# Import the module
module = importlib.import_module(f'.{tool_name}', package='tools')
# Look for tool_instance in the module
if hasattr(module, 'tool_instance'):
tool = module.tool_instance
if isinstance(tool, BaseTool):
self.tools[tool.name] = tool
loaded_count += 1
print(f" ✓ Loaded: {tool.name}")
except Exception as e:
print(f" ✗ Failed to load {tool_name}: {e}")
return loaded_count
def get_tools(self) -> List[Dict]:
"""Get tool definitions in LLM-compatible format (for API calls)"""
return [tool.get_tool_definition() for tool in self.tools.values()]
def get_tool_definitions(self) -> List[MCPTool]:
"""Get tool definitions in MCP format"""
mcp_tools = []
for tool in self.tools.values():
tool_def = tool.get_tool_definition()
# Convert to MCP format
mcp_tool = MCPTool(
name=tool_def["name"],
description=tool_def["description"],
inputSchema=tool_def["inputSchema"]
)
mcp_tools.append(mcp_tool)
return mcp_tools
async def execute_tool(self, tool_name: str, arguments: Dict) -> List[Dict]:
"""Execute a tool by name"""
if tool_name not in self.tools:
raise KeyError(f"Tool '{tool_name}' not found")
tool = self.tools[tool_name]
return await tool.execute(arguments)