Spaces:
Paused
Paused
| """Security agent implementation.""" | |
| from __future__ import annotations | |
| import json | |
| import logging | |
| from typing import Any | |
| from hermes.agents.base.agent import BaseAgent | |
| from hermes.core.types import AgentStrategy, ToolCall, ToolResult | |
| from hermes.tools.base.registry import tool_registry | |
| logger = logging.getLogger(__name__) | |
| class SecurityAgent(BaseAgent): | |
| """Agent specialized in security analysis and vulnerability scanning.""" | |
| def __init__(self, llm_provider: Any | None = None) -> None: | |
| super().__init__( | |
| agent_type="security", | |
| strategy=AgentStrategy.REACT, | |
| tools=["security_scanner", "github_repo_reader", "file_reader"], | |
| llm_provider=llm_provider, | |
| ) | |
| async def plan(self, task: str) -> list[str]: | |
| """Create security analysis plan.""" | |
| return [ | |
| f"Scan for security vulnerabilities: {task}", | |
| "Review code for secrets and sensitive data exposure", | |
| "Check dependency security and configuration", | |
| "Analyze authentication and authorization patterns", | |
| "Compile security assessment report", | |
| ] | |
| async def think(self, task: str, observations: list[str]) -> dict[str, Any]: | |
| """Reason about security approach using LLM.""" | |
| observations_text = "\n".join(f"- {obs[:500]}" for obs in observations) if observations else "None yet." | |
| prompt = f"""You are a security analysis agent. Your task: {task} | |
| Available tools (CHOOSE ONE): | |
| 1. security_scanner — Scan code for vulnerabilities and secrets | |
| Required: {{"action": "scan_directory", "path": "."}} | |
| Actions: scan_code (scan single file), scan_directory (scan folder) | |
| Optional: {{"code": "code string to scan"}} | |
| 2. github_repo_reader — Read GitHub repository files | |
| Required: {{"action": "get_readme", "owner": "owner_name", "repo": "repo_name"}} | |
| Actions: get_readme, list_files, read_file, get_repo | |
| 3. file_reader — Read a local file | |
| Required: {{"action": "read", "path": "file/path.txt"}} | |
| Previous observations: | |
| {observations_text} | |
| What should be your NEXT action? Choose the most appropriate tool. | |
| Respond in JSON format ONLY: | |
| {{"reasoning": "why this tool", "tool": "tool_name", "arguments": {{"key": "value"}}, "done": false}} | |
| Rules: | |
| - Include ALL required arguments for the tool you choose | |
| - If you have enough information, set "done": true and "tool": "none" | |
| - Do NOT make up tool names — use ONLY the 3 tools listed above""" | |
| response = await self._call_llm([{"role": "user", "content": prompt}]) | |
| parsed = self._parse_json_response(response) | |
| if parsed and "tool" in parsed: | |
| parsed.setdefault("reasoning", "") | |
| parsed.setdefault("arguments", {}) | |
| parsed.setdefault("done", False) | |
| return parsed | |
| # Fallback: default to security scan | |
| return { | |
| "reasoning": f"Starting security scan of current directory for: {task}", | |
| "tool": "security_scanner", | |
| "arguments": {"action": "scan_directory", "path": "."}, | |
| "done": False, | |
| } | |
| async def act(self, thought: dict[str, Any]) -> ToolCall: | |
| """Execute security action based on LLM decision.""" | |
| tool_name = thought.get("tool", "security_scanner") | |
| arguments = thought.get("arguments", {}) | |
| valid_tools = ["security_scanner", "github_repo_reader", "file_reader"] | |
| if tool_name not in valid_tools: | |
| tool_name = "security_scanner" | |
| # Ensure required arguments for each tool | |
| if tool_name == "security_scanner": | |
| if "action" not in arguments: | |
| arguments["action"] = "scan_directory" | |
| if "path" not in arguments: | |
| arguments["path"] = "." | |
| elif tool_name == "github_repo_reader": | |
| if "action" not in arguments: | |
| arguments["action"] = "get_readme" | |
| if "owner" not in arguments: | |
| arguments["owner"] = "" | |
| if "repo" not in arguments: | |
| arguments["repo"] = "" | |
| elif tool_name == "file_reader": | |
| if "action" not in arguments: | |
| arguments["action"] = "read" | |
| if "path" not in arguments: | |
| arguments["path"] = "README.md" | |
| return ToolCall(tool_name=tool_name, arguments=arguments) | |
| async def observe(self, result: ToolResult) -> str: | |
| """Observe security results using LLM to extract key findings.""" | |
| if not hasattr(result, "success") or not result.success: | |
| return f"Tool execution failed: {result}" | |
| output = result.output if hasattr(result, "output") else str(result) | |
| output_text = json.dumps(output, default=str)[:3000] if not isinstance(output, str) else output[:3000] | |
| prompt = f"""Extract key security findings from this scan result. | |
| Focus on: secrets, vulnerabilities, risk level, critical issues. | |
| Provide a concise summary (2-3 sentences max). | |
| Tool result: | |
| {output_text} | |
| Security findings:""" | |
| response = await self._call_llm([{"role": "user", "content": prompt}]) | |
| if response and not response.startswith("[LLM unavailable"): | |
| return response.strip() | |
| # Fallback: parse structure | |
| if isinstance(output, dict): | |
| summary = output.get("summary", {}) | |
| if isinstance(summary, dict): | |
| secrets = summary.get("total_secrets", 0) | |
| vulns = summary.get("total_vulnerabilities", 0) | |
| risk = summary.get("risk_level", "unknown") | |
| return f"Security scan: {secrets} secrets, {vulns} vulnerabilities, risk level: {risk}" | |
| return f"Security result: {str(output)[:800]}" | |
| return f"Got result: {str(output)[:800]}" | |
| async def synthesize(self, task: str) -> str: | |
| """Synthesize security findings using LLM.""" | |
| observations = self.state.observations | |
| if not observations: | |
| return f"Security analysis completed for: {task}. No vulnerabilities found." | |
| observations_text = "\n\n".join(f"Finding {i+1}: {obs}" for i, obs in enumerate(observations[:10])) | |
| prompt = f"""You are a security analysis agent synthesizing findings for: | |
| Task: {task} | |
| Security findings: | |
| {observations_text} | |
| Please synthesize these into a comprehensive security assessment: | |
| 1. Critical vulnerabilities found | |
| 2. Overall risk assessment | |
| 3. Recommendations for remediation | |
| 4. Security best practices to implement | |
| Assessment:""" | |
| response = await self._call_llm([{"role": "user", "content": prompt}]) | |
| if response and not response.startswith("[LLM unavailable"): | |
| return response.strip() | |
| summary = f"Security Assessment for: {task}\n\n" | |
| summary += f"Analyzed {len(observations)} security findings.\n\n" | |
| for i, obs in enumerate(observations[:5], 1): | |
| summary += f"Finding {i}: {obs[:300]}\n\n" | |
| return summary | |
| async def scan_repo(self, owner: str, repo: str) -> dict[str, Any]: | |
| """Scan a GitHub repository for security issues.""" | |
| tool = tool_registry.get("github_repo_reader") | |
| if tool: | |
| result = await tool.execute(action="get_readme", owner=owner, repo=repo) | |
| return result if isinstance(result, dict) else {"result": str(result)} | |
| return {"error": "GitHub tool not available"} | |
| async def scan_path(self, path: str) -> dict[str, Any]: | |
| """Scan a local path for security issues.""" | |
| tool = tool_registry.get("security_scanner") | |
| if tool: | |
| result = await tool.execute(action="scan_directory", path=path) | |
| return result if isinstance(result, dict) else {"result": str(result)} | |
| return {"error": "Security scanner not available"} | |