150 lines
5.0 KiB
Python
150 lines
5.0 KiB
Python
from typing import Any
|
|
|
|
from strix.tools.registry import register_tool
|
|
|
|
|
|
def _validate_root_agent(agent_state: Any) -> dict[str, Any] | None:
|
|
if agent_state and hasattr(agent_state, "parent_id") and agent_state.parent_id is not None:
|
|
return {
|
|
"success": False,
|
|
"error": "finish_scan_wrong_agent",
|
|
"message": "This tool can only be used by the root/main agent",
|
|
"suggestion": "If you are a subagent, use agent_finish from agents_graph tool instead",
|
|
}
|
|
return None
|
|
|
|
|
|
def _check_active_agents(agent_state: Any = None) -> dict[str, Any] | None:
|
|
try:
|
|
from strix.tools.agents_graph.agents_graph_actions import _agent_graph
|
|
|
|
if agent_state and agent_state.agent_id:
|
|
current_agent_id = agent_state.agent_id
|
|
else:
|
|
return None
|
|
|
|
active_agents = []
|
|
stopping_agents = []
|
|
|
|
for agent_id, node in _agent_graph["nodes"].items():
|
|
if agent_id == current_agent_id:
|
|
continue
|
|
|
|
status = node.get("status", "unknown")
|
|
if status == "running":
|
|
active_agents.append(
|
|
{
|
|
"id": agent_id,
|
|
"name": node.get("name", "Unknown"),
|
|
"task": node.get("task", "Unknown task")[:300],
|
|
"status": status,
|
|
}
|
|
)
|
|
elif status == "stopping":
|
|
stopping_agents.append(
|
|
{
|
|
"id": agent_id,
|
|
"name": node.get("name", "Unknown"),
|
|
"task": node.get("task", "Unknown task")[:300],
|
|
"status": status,
|
|
}
|
|
)
|
|
|
|
if active_agents or stopping_agents:
|
|
response: dict[str, Any] = {
|
|
"success": False,
|
|
"error": "agents_still_active",
|
|
"message": "Cannot finish scan: agents are still active",
|
|
}
|
|
|
|
if active_agents:
|
|
response["active_agents"] = active_agents
|
|
|
|
if stopping_agents:
|
|
response["stopping_agents"] = stopping_agents
|
|
|
|
response["suggestions"] = [
|
|
"Use wait_for_message to wait for all agents to complete",
|
|
"Use send_message_to_agent if you need agents to complete immediately",
|
|
"Check agent_status to see current agent states",
|
|
]
|
|
|
|
response["total_active"] = len(active_agents) + len(stopping_agents)
|
|
|
|
return response
|
|
|
|
except ImportError:
|
|
pass
|
|
except Exception:
|
|
import logging
|
|
|
|
logging.exception("Error checking active agents")
|
|
|
|
return None
|
|
|
|
|
|
@register_tool(sandbox_execution=False)
|
|
def finish_scan(
|
|
executive_summary: str,
|
|
methodology: str,
|
|
technical_analysis: str,
|
|
recommendations: str,
|
|
agent_state: Any = None,
|
|
) -> dict[str, Any]:
|
|
validation_error = _validate_root_agent(agent_state)
|
|
if validation_error:
|
|
return validation_error
|
|
|
|
active_agents_error = _check_active_agents(agent_state)
|
|
if active_agents_error:
|
|
return active_agents_error
|
|
|
|
validation_errors = []
|
|
|
|
if not executive_summary or not executive_summary.strip():
|
|
validation_errors.append("Executive summary cannot be empty")
|
|
if not methodology or not methodology.strip():
|
|
validation_errors.append("Methodology cannot be empty")
|
|
if not technical_analysis or not technical_analysis.strip():
|
|
validation_errors.append("Technical analysis cannot be empty")
|
|
if not recommendations or not recommendations.strip():
|
|
validation_errors.append("Recommendations cannot be empty")
|
|
|
|
if validation_errors:
|
|
return {"success": False, "message": "Validation failed", "errors": validation_errors}
|
|
|
|
try:
|
|
from strix.telemetry.tracer import get_global_tracer
|
|
|
|
tracer = get_global_tracer()
|
|
if tracer:
|
|
tracer.update_scan_final_fields(
|
|
executive_summary=executive_summary.strip(),
|
|
methodology=methodology.strip(),
|
|
technical_analysis=technical_analysis.strip(),
|
|
recommendations=recommendations.strip(),
|
|
)
|
|
|
|
vulnerability_count = len(tracer.vulnerability_reports)
|
|
|
|
return {
|
|
"success": True,
|
|
"scan_completed": True,
|
|
"message": "Scan completed successfully",
|
|
"vulnerabilities_found": vulnerability_count,
|
|
}
|
|
|
|
import logging
|
|
|
|
logging.warning("Current tracer not available - scan results not stored")
|
|
|
|
except (ImportError, AttributeError) as e:
|
|
return {"success": False, "message": f"Failed to complete scan: {e!s}"}
|
|
else:
|
|
return {
|
|
"success": True,
|
|
"scan_completed": True,
|
|
"message": "Scan completed (not persisted)",
|
|
"warning": "Results could not be persisted - tracer unavailable",
|
|
}
|