609 lines
22 KiB
Python
609 lines
22 KiB
Python
|
"""
|
||
|
|
AUTARCH Tool System
|
||
|
|
Defines tools that the agent can use to interact with the environment
|
||
|
|
"""
|
||
|
|
|
||
|
|
import os
|
||
|
|
import subprocess
|
||
|
|
import json
|
||
|
|
from typing import Callable, Dict, List, Any, Optional
|
||
|
|
from dataclasses import dataclass, field
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
from .banner import Colors
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass
|
||
|
|
class ToolParameter:
|
||
|
|
"""Definition of a tool parameter."""
|
||
|
|
name: str
|
||
|
|
description: str
|
||
|
|
type: str = "string"
|
||
|
|
required: bool = True
|
||
|
|
default: Any = None
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass
|
||
|
|
class Tool:
|
||
|
|
"""Definition of an agent tool."""
|
||
|
|
name: str
|
||
|
|
description: str
|
||
|
|
function: Callable
|
||
|
|
parameters: List[ToolParameter] = field(default_factory=list)
|
||
|
|
category: str = "general"
|
||
|
|
|
||
|
|
def to_schema(self) -> Dict[str, Any]:
|
||
|
|
"""Convert tool to JSON schema for LLM."""
|
||
|
|
properties = {}
|
||
|
|
required = []
|
||
|
|
|
||
|
|
for param in self.parameters:
|
||
|
|
properties[param.name] = {
|
||
|
|
"type": param.type,
|
||
|
|
"description": param.description
|
||
|
|
}
|
||
|
|
if param.required:
|
||
|
|
required.append(param.name)
|
||
|
|
|
||
|
|
return {
|
||
|
|
"name": self.name,
|
||
|
|
"description": self.description,
|
||
|
|
"parameters": {
|
||
|
|
"type": "object",
|
||
|
|
"properties": properties,
|
||
|
|
"required": required
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
def execute(self, **kwargs) -> Dict[str, Any]:
|
||
|
|
"""Execute the tool with given parameters.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
Dict with 'success' bool and 'result' or 'error' string.
|
||
|
|
"""
|
||
|
|
try:
|
||
|
|
result = self.function(**kwargs)
|
||
|
|
return {"success": True, "result": result}
|
||
|
|
except Exception as e:
|
||
|
|
return {"success": False, "error": str(e)}
|
||
|
|
|
||
|
|
|
||
|
|
class ToolRegistry:
|
||
|
|
"""Registry for managing available tools."""
|
||
|
|
|
||
|
|
def __init__(self):
|
||
|
|
self._tools: Dict[str, Tool] = {}
|
||
|
|
self._register_builtin_tools()
|
||
|
|
|
||
|
|
def register(self, tool: Tool):
|
||
|
|
"""Register a tool."""
|
||
|
|
self._tools[tool.name] = tool
|
||
|
|
|
||
|
|
def unregister(self, name: str):
|
||
|
|
"""Unregister a tool by name."""
|
||
|
|
if name in self._tools:
|
||
|
|
del self._tools[name]
|
||
|
|
|
||
|
|
def get(self, name: str) -> Optional[Tool]:
|
||
|
|
"""Get a tool by name."""
|
||
|
|
return self._tools.get(name)
|
||
|
|
|
||
|
|
def list_tools(self) -> List[Tool]:
|
||
|
|
"""List all registered tools."""
|
||
|
|
return list(self._tools.values())
|
||
|
|
|
||
|
|
def get_tools_schema(self) -> List[Dict[str, Any]]:
|
||
|
|
"""Get JSON schema for all tools."""
|
||
|
|
return [tool.to_schema() for tool in self._tools.values()]
|
||
|
|
|
||
|
|
def get_tools_prompt(self) -> str:
|
||
|
|
"""Generate a tools description for the LLM prompt."""
|
||
|
|
lines = ["Available tools:"]
|
||
|
|
for tool in self._tools.values():
|
||
|
|
lines.append(f"\n## {tool.name}")
|
||
|
|
lines.append(f"Description: {tool.description}")
|
||
|
|
if tool.parameters:
|
||
|
|
lines.append("Parameters:")
|
||
|
|
for param in tool.parameters:
|
||
|
|
req = "(required)" if param.required else "(optional)"
|
||
|
|
lines.append(f" - {param.name} [{param.type}] {req}: {param.description}")
|
||
|
|
return "\n".join(lines)
|
||
|
|
|
||
|
|
def execute(self, tool_name: str, **kwargs) -> Dict[str, Any]:
|
||
|
|
"""Execute a tool by name.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
tool_name: Name of the tool to execute.
|
||
|
|
**kwargs: Parameters to pass to the tool.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
Dict with execution result.
|
||
|
|
"""
|
||
|
|
tool = self.get(tool_name)
|
||
|
|
if not tool:
|
||
|
|
return {"success": False, "error": f"Tool '{tool_name}' not found"}
|
||
|
|
return tool.execute(**kwargs)
|
||
|
|
|
||
|
|
def _register_builtin_tools(self):
|
||
|
|
"""Register built-in tools."""
|
||
|
|
|
||
|
|
# Shell command execution
|
||
|
|
self.register(Tool(
|
||
|
|
name="shell",
|
||
|
|
description="Execute a shell command and return the output. Use for system operations, running scripts, or gathering system information.",
|
||
|
|
function=self._tool_shell,
|
||
|
|
parameters=[
|
||
|
|
ToolParameter("command", "The shell command to execute", "string", True),
|
||
|
|
ToolParameter("timeout", "Timeout in seconds (default 30)", "integer", False, 30),
|
||
|
|
],
|
||
|
|
category="system"
|
||
|
|
))
|
||
|
|
|
||
|
|
# Read file
|
||
|
|
self.register(Tool(
|
||
|
|
name="read_file",
|
||
|
|
description="Read the contents of a file. Use to examine files, configs, or source code.",
|
||
|
|
function=self._tool_read_file,
|
||
|
|
parameters=[
|
||
|
|
ToolParameter("path", "Path to the file to read", "string", True),
|
||
|
|
ToolParameter("max_lines", "Maximum number of lines to read (default all)", "integer", False),
|
||
|
|
],
|
||
|
|
category="filesystem"
|
||
|
|
))
|
||
|
|
|
||
|
|
# Write file
|
||
|
|
self.register(Tool(
|
||
|
|
name="write_file",
|
||
|
|
description="Write content to a file. Creates the file if it doesn't exist, overwrites if it does.",
|
||
|
|
function=self._tool_write_file,
|
||
|
|
parameters=[
|
||
|
|
ToolParameter("path", "Path to the file to write", "string", True),
|
||
|
|
ToolParameter("content", "Content to write to the file", "string", True),
|
||
|
|
],
|
||
|
|
category="filesystem"
|
||
|
|
))
|
||
|
|
|
||
|
|
# List directory
|
||
|
|
self.register(Tool(
|
||
|
|
name="list_dir",
|
||
|
|
description="List contents of a directory. Use to explore filesystem structure.",
|
||
|
|
function=self._tool_list_dir,
|
||
|
|
parameters=[
|
||
|
|
ToolParameter("path", "Path to the directory (default: current)", "string", False, "."),
|
||
|
|
ToolParameter("show_hidden", "Include hidden files (default: false)", "boolean", False, False),
|
||
|
|
],
|
||
|
|
category="filesystem"
|
||
|
|
))
|
||
|
|
|
||
|
|
# Search files
|
||
|
|
self.register(Tool(
|
||
|
|
name="search_files",
|
||
|
|
description="Search for files matching a pattern. Use to find specific files.",
|
||
|
|
function=self._tool_search_files,
|
||
|
|
parameters=[
|
||
|
|
ToolParameter("pattern", "Glob pattern to match (e.g., '*.py', '**/*.txt')", "string", True),
|
||
|
|
ToolParameter("path", "Starting directory (default: current)", "string", False, "."),
|
||
|
|
],
|
||
|
|
category="filesystem"
|
||
|
|
))
|
||
|
|
|
||
|
|
# Search in files (grep)
|
||
|
|
self.register(Tool(
|
||
|
|
name="search_content",
|
||
|
|
description="Search for text content within files. Use to find specific code or text.",
|
||
|
|
function=self._tool_search_content,
|
||
|
|
parameters=[
|
||
|
|
ToolParameter("pattern", "Text or regex pattern to search for", "string", True),
|
||
|
|
ToolParameter("path", "File or directory to search in", "string", False, "."),
|
||
|
|
ToolParameter("file_pattern", "Glob pattern for files to search (e.g., '*.py')", "string", False),
|
||
|
|
],
|
||
|
|
category="filesystem"
|
||
|
|
))
|
||
|
|
|
||
|
|
# Task complete
|
||
|
|
self.register(Tool(
|
||
|
|
name="task_complete",
|
||
|
|
description="Mark the current task as complete. Use when you have fully accomplished the goal.",
|
||
|
|
function=self._tool_task_complete,
|
||
|
|
parameters=[
|
||
|
|
ToolParameter("summary", "Summary of what was accomplished", "string", True),
|
||
|
|
],
|
||
|
|
category="control"
|
||
|
|
))
|
||
|
|
|
||
|
|
# Ask user
|
||
|
|
self.register(Tool(
|
||
|
|
name="ask_user",
|
||
|
|
description="Ask the user a question when you need clarification or input.",
|
||
|
|
function=self._tool_ask_user,
|
||
|
|
parameters=[
|
||
|
|
ToolParameter("question", "The question to ask the user", "string", True),
|
||
|
|
],
|
||
|
|
category="interaction"
|
||
|
|
))
|
||
|
|
|
||
|
|
# Metasploit tools
|
||
|
|
self.register(Tool(
|
||
|
|
name="msf_connect",
|
||
|
|
description="Connect to Metasploit RPC. Required before using other MSF tools.",
|
||
|
|
function=self._tool_msf_connect,
|
||
|
|
parameters=[
|
||
|
|
ToolParameter("password", "MSF RPC password (uses saved if not provided)", "string", False),
|
||
|
|
],
|
||
|
|
category="msf"
|
||
|
|
))
|
||
|
|
|
||
|
|
self.register(Tool(
|
||
|
|
name="msf_search",
|
||
|
|
description="Search for Metasploit modules by keyword.",
|
||
|
|
function=self._tool_msf_search,
|
||
|
|
parameters=[
|
||
|
|
ToolParameter("query", "Search query (e.g., 'smb', 'apache', 'cve:2021')", "string", True),
|
||
|
|
],
|
||
|
|
category="msf"
|
||
|
|
))
|
||
|
|
|
||
|
|
self.register(Tool(
|
||
|
|
name="msf_module_info",
|
||
|
|
description="Get detailed information about a Metasploit module.",
|
||
|
|
function=self._tool_msf_module_info,
|
||
|
|
parameters=[
|
||
|
|
ToolParameter("module_type", "Module type: exploit, auxiliary, post, payload", "string", True),
|
||
|
|
ToolParameter("module_name", "Module name (e.g., 'windows/smb/ms17_010_eternalblue')", "string", True),
|
||
|
|
],
|
||
|
|
category="msf"
|
||
|
|
))
|
||
|
|
|
||
|
|
self.register(Tool(
|
||
|
|
name="msf_module_options",
|
||
|
|
description="Get available options for a Metasploit module.",
|
||
|
|
function=self._tool_msf_module_options,
|
||
|
|
parameters=[
|
||
|
|
ToolParameter("module_type", "Module type: exploit, auxiliary, post, payload", "string", True),
|
||
|
|
ToolParameter("module_name", "Module name", "string", True),
|
||
|
|
],
|
||
|
|
category="msf"
|
||
|
|
))
|
||
|
|
|
||
|
|
self.register(Tool(
|
||
|
|
name="msf_execute",
|
||
|
|
description="Execute a Metasploit module with specified options.",
|
||
|
|
function=self._tool_msf_execute,
|
||
|
|
parameters=[
|
||
|
|
ToolParameter("module_type", "Module type: exploit, auxiliary, post", "string", True),
|
||
|
|
ToolParameter("module_name", "Module name", "string", True),
|
||
|
|
ToolParameter("options", "JSON object of module options (e.g., {\"RHOSTS\": \"192.168.1.1\"})", "string", True),
|
||
|
|
],
|
||
|
|
category="msf"
|
||
|
|
))
|
||
|
|
|
||
|
|
self.register(Tool(
|
||
|
|
name="msf_sessions",
|
||
|
|
description="List active Metasploit sessions.",
|
||
|
|
function=self._tool_msf_sessions,
|
||
|
|
parameters=[],
|
||
|
|
category="msf"
|
||
|
|
))
|
||
|
|
|
||
|
|
self.register(Tool(
|
||
|
|
name="msf_session_command",
|
||
|
|
description="Execute a command in a Metasploit session.",
|
||
|
|
function=self._tool_msf_session_command,
|
||
|
|
parameters=[
|
||
|
|
ToolParameter("session_id", "Session ID", "string", True),
|
||
|
|
ToolParameter("command", "Command to execute", "string", True),
|
||
|
|
],
|
||
|
|
category="msf"
|
||
|
|
))
|
||
|
|
|
||
|
|
self.register(Tool(
|
||
|
|
name="msf_console",
|
||
|
|
description="Run a command in the Metasploit console.",
|
||
|
|
function=self._tool_msf_console,
|
||
|
|
parameters=[
|
||
|
|
ToolParameter("command", "Console command to run", "string", True),
|
||
|
|
],
|
||
|
|
category="msf"
|
||
|
|
))
|
||
|
|
|
||
|
|
# Built-in tool implementations
|
||
|
|
|
||
|
|
def _tool_shell(self, command: str, timeout: int = 30) -> str:
|
||
|
|
"""Execute a shell command."""
|
||
|
|
try:
|
||
|
|
result = subprocess.run(
|
||
|
|
command,
|
||
|
|
shell=True,
|
||
|
|
capture_output=True,
|
||
|
|
text=True,
|
||
|
|
timeout=timeout
|
||
|
|
)
|
||
|
|
output = result.stdout
|
||
|
|
if result.stderr:
|
||
|
|
output += f"\n[stderr]: {result.stderr}"
|
||
|
|
if result.returncode != 0:
|
||
|
|
output += f"\n[exit code]: {result.returncode}"
|
||
|
|
return output.strip() or "[no output]"
|
||
|
|
except subprocess.TimeoutExpired:
|
||
|
|
return f"[error]: Command timed out after {timeout} seconds"
|
||
|
|
except Exception as e:
|
||
|
|
return f"[error]: {str(e)}"
|
||
|
|
|
||
|
|
def _tool_read_file(self, path: str, max_lines: int = None) -> str:
|
||
|
|
"""Read a file's contents."""
|
||
|
|
path = Path(path).expanduser()
|
||
|
|
if not path.exists():
|
||
|
|
raise FileNotFoundError(f"File not found: {path}")
|
||
|
|
if not path.is_file():
|
||
|
|
raise ValueError(f"Not a file: {path}")
|
||
|
|
|
||
|
|
with open(path, 'r', errors='replace') as f:
|
||
|
|
if max_lines:
|
||
|
|
lines = []
|
||
|
|
for i, line in enumerate(f):
|
||
|
|
if i >= max_lines:
|
||
|
|
lines.append(f"... [{path.stat().st_size} bytes total, truncated at {max_lines} lines]")
|
||
|
|
break
|
||
|
|
lines.append(line.rstrip())
|
||
|
|
return '\n'.join(lines)
|
||
|
|
else:
|
||
|
|
return f.read()
|
||
|
|
|
||
|
|
def _tool_write_file(self, path: str, content: str) -> str:
|
||
|
|
"""Write content to a file."""
|
||
|
|
path = Path(path).expanduser()
|
||
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
||
|
|
with open(path, 'w') as f:
|
||
|
|
f.write(content)
|
||
|
|
return f"Successfully wrote {len(content)} bytes to {path}"
|
||
|
|
|
||
|
|
def _tool_list_dir(self, path: str = ".", show_hidden: bool = False) -> str:
|
||
|
|
"""List directory contents."""
|
||
|
|
path = Path(path).expanduser()
|
||
|
|
if not path.exists():
|
||
|
|
raise FileNotFoundError(f"Directory not found: {path}")
|
||
|
|
if not path.is_dir():
|
||
|
|
raise ValueError(f"Not a directory: {path}")
|
||
|
|
|
||
|
|
entries = []
|
||
|
|
for entry in sorted(path.iterdir()):
|
||
|
|
if not show_hidden and entry.name.startswith('.'):
|
||
|
|
continue
|
||
|
|
prefix = "d " if entry.is_dir() else "f "
|
||
|
|
entries.append(f"{prefix}{entry.name}")
|
||
|
|
|
||
|
|
return '\n'.join(entries) if entries else "[empty directory]"
|
||
|
|
|
||
|
|
def _tool_search_files(self, pattern: str, path: str = ".") -> str:
|
||
|
|
"""Search for files matching a pattern."""
|
||
|
|
path = Path(path).expanduser()
|
||
|
|
matches = list(path.glob(pattern))
|
||
|
|
|
||
|
|
if not matches:
|
||
|
|
return f"No files matching '{pattern}'"
|
||
|
|
|
||
|
|
result = []
|
||
|
|
for match in matches[:50]: # Limit results
|
||
|
|
result.append(str(match))
|
||
|
|
|
||
|
|
if len(matches) > 50:
|
||
|
|
result.append(f"... and {len(matches) - 50} more")
|
||
|
|
|
||
|
|
return '\n'.join(result)
|
||
|
|
|
||
|
|
def _tool_search_content(self, pattern: str, path: str = ".", file_pattern: str = None) -> str:
|
||
|
|
"""Search for content in files."""
|
||
|
|
try:
|
||
|
|
cmd = f"grep -rn '{pattern}' {path}"
|
||
|
|
if file_pattern:
|
||
|
|
cmd = f"grep -rn --include='{file_pattern}' '{pattern}' {path}"
|
||
|
|
|
||
|
|
result = subprocess.run(
|
||
|
|
cmd,
|
||
|
|
shell=True,
|
||
|
|
capture_output=True,
|
||
|
|
text=True,
|
||
|
|
timeout=30
|
||
|
|
)
|
||
|
|
|
||
|
|
output = result.stdout.strip()
|
||
|
|
if not output:
|
||
|
|
return f"No matches found for '{pattern}'"
|
||
|
|
|
||
|
|
# Limit output
|
||
|
|
lines = output.split('\n')
|
||
|
|
if len(lines) > 30:
|
||
|
|
return '\n'.join(lines[:30]) + f"\n... and {len(lines) - 30} more matches"
|
||
|
|
return output
|
||
|
|
|
||
|
|
except subprocess.TimeoutExpired:
|
||
|
|
return "[error]: Search timed out"
|
||
|
|
except Exception as e:
|
||
|
|
return f"[error]: {str(e)}"
|
||
|
|
|
||
|
|
def _tool_task_complete(self, summary: str) -> str:
|
||
|
|
"""Mark task as complete - this is a control signal."""
|
||
|
|
return f"__TASK_COMPLETE__:{summary}"
|
||
|
|
|
||
|
|
def _tool_ask_user(self, question: str) -> str:
|
||
|
|
"""Ask user a question - handled by agent loop."""
|
||
|
|
return f"__ASK_USER__:{question}"
|
||
|
|
|
||
|
|
# Metasploit tool implementations
|
||
|
|
|
||
|
|
def _tool_msf_connect(self, password: str = None) -> str:
|
||
|
|
"""Connect to Metasploit RPC."""
|
||
|
|
from .msf import get_msf_manager, MSFError
|
||
|
|
|
||
|
|
msf = get_msf_manager()
|
||
|
|
try:
|
||
|
|
msf.connect(password)
|
||
|
|
version = msf.rpc.get_version()
|
||
|
|
return f"Connected to Metasploit {version.get('version', 'Unknown')}"
|
||
|
|
except MSFError as e:
|
||
|
|
return f"[error]: {e}"
|
||
|
|
|
||
|
|
def _tool_msf_search(self, query: str) -> str:
|
||
|
|
"""Search for Metasploit modules."""
|
||
|
|
from .msf import get_msf_manager, MSFError
|
||
|
|
|
||
|
|
msf = get_msf_manager()
|
||
|
|
if not msf.is_connected:
|
||
|
|
return "[error]: Not connected to Metasploit. Use msf_connect first."
|
||
|
|
|
||
|
|
try:
|
||
|
|
results = msf.rpc.search_modules(query)
|
||
|
|
if not results:
|
||
|
|
return f"No modules found matching '{query}'"
|
||
|
|
|
||
|
|
output = []
|
||
|
|
for i, mod in enumerate(results[:20]): # Limit to 20 results
|
||
|
|
if isinstance(mod, dict):
|
||
|
|
name = mod.get('fullname', mod.get('name', 'Unknown'))
|
||
|
|
desc = mod.get('description', '')[:60]
|
||
|
|
output.append(f"{name}\n {desc}")
|
||
|
|
else:
|
||
|
|
output.append(str(mod))
|
||
|
|
|
||
|
|
if len(results) > 20:
|
||
|
|
output.append(f"\n... and {len(results) - 20} more results")
|
||
|
|
|
||
|
|
return '\n'.join(output)
|
||
|
|
except MSFError as e:
|
||
|
|
return f"[error]: {e}"
|
||
|
|
|
||
|
|
def _tool_msf_module_info(self, module_type: str, module_name: str) -> str:
|
||
|
|
"""Get module information."""
|
||
|
|
from .msf import get_msf_manager, MSFError
|
||
|
|
|
||
|
|
msf = get_msf_manager()
|
||
|
|
if not msf.is_connected:
|
||
|
|
return "[error]: Not connected to Metasploit. Use msf_connect first."
|
||
|
|
|
||
|
|
try:
|
||
|
|
info = msf.rpc.get_module_info(module_type, module_name)
|
||
|
|
output = [
|
||
|
|
f"Name: {info.name}",
|
||
|
|
f"Type: {info.type}",
|
||
|
|
f"Rank: {info.rank}",
|
||
|
|
f"Description: {info.description[:200]}..." if len(info.description) > 200 else f"Description: {info.description}",
|
||
|
|
]
|
||
|
|
if info.author:
|
||
|
|
output.append(f"Authors: {', '.join(info.author[:3])}")
|
||
|
|
return '\n'.join(output)
|
||
|
|
except MSFError as e:
|
||
|
|
return f"[error]: {e}"
|
||
|
|
|
||
|
|
def _tool_msf_module_options(self, module_type: str, module_name: str) -> str:
|
||
|
|
"""Get module options."""
|
||
|
|
from .msf import get_msf_manager, MSFError
|
||
|
|
|
||
|
|
msf = get_msf_manager()
|
||
|
|
if not msf.is_connected:
|
||
|
|
return "[error]: Not connected to Metasploit. Use msf_connect first."
|
||
|
|
|
||
|
|
try:
|
||
|
|
options = msf.rpc.get_module_options(module_type, module_name)
|
||
|
|
output = []
|
||
|
|
for name, details in options.items():
|
||
|
|
if isinstance(details, dict):
|
||
|
|
required = "*" if details.get('required', False) else ""
|
||
|
|
default = details.get('default', '')
|
||
|
|
desc = details.get('desc', '')[:50]
|
||
|
|
output.append(f"{name}{required}: {desc} [default: {default}]")
|
||
|
|
else:
|
||
|
|
output.append(f"{name}: {details}")
|
||
|
|
return '\n'.join(output) if output else "No options available"
|
||
|
|
except MSFError as e:
|
||
|
|
return f"[error]: {e}"
|
||
|
|
|
||
|
|
def _tool_msf_execute(self, module_type: str, module_name: str, options: str) -> str:
|
||
|
|
"""Execute a Metasploit module."""
|
||
|
|
from .msf import get_msf_manager, MSFError
|
||
|
|
|
||
|
|
msf = get_msf_manager()
|
||
|
|
if not msf.is_connected:
|
||
|
|
return "[error]: Not connected to Metasploit. Use msf_connect first."
|
||
|
|
|
||
|
|
try:
|
||
|
|
opts = json.loads(options) if isinstance(options, str) else options
|
||
|
|
except json.JSONDecodeError:
|
||
|
|
return "[error]: Invalid JSON in options parameter"
|
||
|
|
|
||
|
|
try:
|
||
|
|
result = msf.rpc.execute_module(module_type, module_name, opts)
|
||
|
|
job_id = result.get('job_id')
|
||
|
|
uuid = result.get('uuid')
|
||
|
|
return f"Module executed. Job ID: {job_id}, UUID: {uuid}"
|
||
|
|
except MSFError as e:
|
||
|
|
return f"[error]: {e}"
|
||
|
|
|
||
|
|
def _tool_msf_sessions(self) -> str:
|
||
|
|
"""List active sessions."""
|
||
|
|
from .msf import get_msf_manager, MSFError
|
||
|
|
|
||
|
|
msf = get_msf_manager()
|
||
|
|
if not msf.is_connected:
|
||
|
|
return "[error]: Not connected to Metasploit. Use msf_connect first."
|
||
|
|
|
||
|
|
try:
|
||
|
|
sessions = msf.rpc.list_sessions()
|
||
|
|
if not sessions:
|
||
|
|
return "No active sessions"
|
||
|
|
|
||
|
|
output = []
|
||
|
|
for sid, info in sessions.items():
|
||
|
|
if isinstance(info, dict):
|
||
|
|
stype = info.get('type', 'Unknown')
|
||
|
|
target = info.get('target_host', 'Unknown')
|
||
|
|
user = info.get('username', '')
|
||
|
|
output.append(f"[{sid}] {stype} - {target} ({user})")
|
||
|
|
else:
|
||
|
|
output.append(f"[{sid}] {info}")
|
||
|
|
return '\n'.join(output)
|
||
|
|
except MSFError as e:
|
||
|
|
return f"[error]: {e}"
|
||
|
|
|
||
|
|
def _tool_msf_session_command(self, session_id: str, command: str) -> str:
|
||
|
|
"""Execute command in a session."""
|
||
|
|
from .msf import get_msf_manager, MSFError
|
||
|
|
|
||
|
|
msf = get_msf_manager()
|
||
|
|
if not msf.is_connected:
|
||
|
|
return "[error]: Not connected to Metasploit. Use msf_connect first."
|
||
|
|
|
||
|
|
try:
|
||
|
|
msf.rpc.session_shell_write(session_id, command)
|
||
|
|
import time
|
||
|
|
time.sleep(1) # Wait for command execution
|
||
|
|
output = msf.rpc.session_shell_read(session_id)
|
||
|
|
return output if output else "[no output]"
|
||
|
|
except MSFError as e:
|
||
|
|
return f"[error]: {e}"
|
||
|
|
|
||
|
|
def _tool_msf_console(self, command: str) -> str:
|
||
|
|
"""Run a console command."""
|
||
|
|
from .msf import get_msf_manager, MSFError
|
||
|
|
|
||
|
|
msf = get_msf_manager()
|
||
|
|
if not msf.is_connected:
|
||
|
|
return "[error]: Not connected to Metasploit. Use msf_connect first."
|
||
|
|
|
||
|
|
try:
|
||
|
|
output = msf.rpc.run_console_command(command)
|
||
|
|
return output if output else "[no output]"
|
||
|
|
except MSFError as e:
|
||
|
|
return f"[error]: {e}"
|
||
|
|
|
||
|
|
|
||
|
|
# Global tool registry
|
||
|
|
_registry: Optional[ToolRegistry] = None
|
||
|
|
|
||
|
|
|
||
|
|
def get_tool_registry() -> ToolRegistry:
|
||
|
|
"""Get the global tool registry."""
|
||
|
|
global _registry
|
||
|
|
if _registry is None:
|
||
|
|
_registry = ToolRegistry()
|
||
|
|
return _registry
|