Autarch Will Control The Internet
This commit is contained in:
DigiJ
675
core/tools.py
Normal file
675
core/tools.py
Normal file
@@ -0,0 +1,675 @@
|
||||
"""
|
||||
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"
|
||||
))
|
||||
|
||||
# Create module
|
||||
self.register(Tool(
|
||||
name="create_module",
|
||||
description="Create a new AUTARCH module. Writes a Python file to the modules/ directory that becomes available in the dashboard.",
|
||||
function=self._tool_create_module,
|
||||
parameters=[
|
||||
ToolParameter("name", "Module filename without .py extension (e.g., port_scanner)", "string", True),
|
||||
ToolParameter("category", "Module category: defense, offense, counter, analyze, osint, or simulate", "string", True),
|
||||
ToolParameter("code", "Complete Python source code for the module", "string", True),
|
||||
],
|
||||
category="development"
|
||||
))
|
||||
|
||||
# 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_create_module(self, name: str, category: str, code: str) -> str:
|
||||
"""Create a new AUTARCH module in the modules/ directory."""
|
||||
import importlib.util as ilu
|
||||
|
||||
valid_categories = ('defense', 'offense', 'counter', 'analyze', 'osint', 'simulate')
|
||||
category = category.lower().strip()
|
||||
if category not in valid_categories:
|
||||
return f"[error]: Invalid category '{category}'. Must be one of: {', '.join(valid_categories)}"
|
||||
|
||||
# Sanitize name
|
||||
name = name.strip().replace(' ', '_').replace('-', '_').lower()
|
||||
if not name.replace('_', '').isalnum():
|
||||
return f"[error]: Invalid module name '{name}'. Use only letters, numbers, and underscores."
|
||||
|
||||
# Check required attributes in source code
|
||||
required = ['DESCRIPTION', 'VERSION', 'CATEGORY', 'def run(']
|
||||
missing = [r for r in required if r not in code]
|
||||
if missing:
|
||||
return f"[error]: Module code is missing required elements: {', '.join(missing)}"
|
||||
|
||||
# Determine modules directory
|
||||
modules_dir = Path(__file__).parent.parent / 'modules'
|
||||
module_path = modules_dir / f'{name}.py'
|
||||
|
||||
if module_path.exists():
|
||||
return f"[error]: Module '{name}' already exists at {module_path}. Choose a different name."
|
||||
|
||||
# Write the module file
|
||||
try:
|
||||
module_path.write_text(code, encoding='utf-8')
|
||||
except Exception as e:
|
||||
return f"[error]: Failed to write module: {e}"
|
||||
|
||||
# Validate by attempting to import
|
||||
try:
|
||||
spec = ilu.spec_from_file_location(name, module_path)
|
||||
mod = ilu.module_from_spec(spec)
|
||||
spec.loader.exec_module(mod)
|
||||
|
||||
# Verify it has run()
|
||||
if not hasattr(mod, 'run'):
|
||||
module_path.unlink()
|
||||
return "[error]: Module loaded but has no run() function. Module deleted."
|
||||
|
||||
except Exception as e:
|
||||
# Import failed — delete the bad module
|
||||
try:
|
||||
module_path.unlink()
|
||||
except Exception:
|
||||
pass
|
||||
return f"[error]: Module failed to import: {e}. Module deleted."
|
||||
|
||||
return f"Module '{name}' created successfully at {module_path}. Category: {category}. It is now available in the dashboard."
|
||||
|
||||
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
|
||||
Reference in New Issue
Block a user