zen-ai-pentest 2.0.0__py3-none-any.whl

agents/exploit_agent.py

@@ -0,0 +1,224 @@
+#!/usr/bin/env python3
+"""
+Exploit Agent - Specializes in developing exploits and payloads
+Part of the Multi-Agent Collaboration System
+Author: SHAdd0WTAka
+"""
+
+import asyncio
+import logging
+from typing import Any, Dict, List
+
+from modules.sql_injection_db import (DBType, SQLInjectionDatabase,
+                                      SQLITechnique)
+
+from .agent_base import AgentMessage, AgentRole, BaseAgent
+
+logger = logging.getLogger("ZenAI.Agents")
+
+
+class ExploitAgent(BaseAgent):
+    """
+    Exploit Agent - Generates exploits, payloads, and attack code
+    Works with Research and Analysis agents for context
+    """
+
+    def __init__(self, name: str, orchestrator=None, zen_orchestrator=None):
+        super().__init__(name, AgentRole.EXPLOIT, orchestrator)
+        self.zen_orchestrator = zen_orchestrator
+        self.sqli_db = SQLInjectionDatabase()
+        self.payload_cache = {}
+
+        # Register handlers
+        self.register_handler("exploit_request", self._handle_exploit_request)
+        self.register_handler("findings", self._handle_findings)
+        self.register_handler("analysis_result", self._handle_analysis_result)
+
+    async def _handle_exploit_request(self, msg: AgentMessage):
+        """Handle requests for exploit development"""
+        vuln_type = msg.context.get("vulnerability_type", "")
+        target_info = msg.context.get("target_info", {})
+
+        logger.info(
+            f"[ExploitAgent:{self.name}] Received exploit request for {vuln_type}"
+        )
+
+        # Generate exploit
+        exploit = await self._develop_exploit(vuln_type, target_info)
+
+        # Send back to requester
+        await self.send_message(
+            content=f"Exploit developed for {vuln_type}",
+            recipient=msg.sender,
+            msg_type="exploit_ready",
+            context={
+                "exploit": exploit,
+                "vulnerability": vuln_type,
+                "author": self.name,
+            },
+        )
+
+        # Also share with all for coordination
+        if exploit.get("severity") in ["critical", "high"]:
+            await self.send_message(
+                content=f"High-severity exploit available: {vuln_type}",
+                recipient="all",
+                msg_type="exploit_notification",
+                priority=3,
+                context={"exploit_summary": exploit.get("summary", "")},
+            )
+
+    async def _handle_findings(self, msg: AgentMessage):
+        """Process findings to identify exploitable vulnerabilities"""
+        findings = msg.context.get("findings", [])
+
+        exploitable = []
+        for finding in findings:
+            if finding.get("type") == "cve":
+                cve_data = finding.get("data", {})
+                severity = getattr(cve_data, "severity", "Unknown")
+                if severity in ["Critical", "High"]:
+                    exploitable.append(finding)
+
+        if exploitable:
+            logger.info(
+                f"[ExploitAgent:{self.name}] Found {len(exploitable)} exploitable vulnerabilities"
+            )
+
+            # Request analysis of these for exploit development
+            await self.send_message(
+                content=f"Requesting analysis of {len(exploitable)} exploitable CVEs",
+                recipient="role:analyst",
+                msg_type="analysis_request",
+                context={"data": exploitable, "analysis_type": "attack_path"},
+            )
+
+    async def _handle_analysis_result(self, msg: AgentMessage):
+        """Handle analysis results to develop targeted exploits"""
+        analysis = msg.context.get("analysis", {})
+
+        # If analysis identified attack paths, develop exploits for them
+        if "paths" in analysis:
+            for path in analysis.get("paths", []):
+                # Develop exploit for each stage
+                pass
+
+    async def _develop_exploit(self, vuln_type: str, target_info: Dict) -> Dict:
+        """Develop an exploit for the given vulnerability"""
+        exploit = {
+            "vulnerability": vuln_type,
+            "target": target_info.get("target", "unknown"),
+            "severity": "medium",
+            "payloads": [],
+            "techniques": [],
+            "mitigation": "",
+            "detection": "",
+            "summary": "",
+        }
+
+        # Use SQL injection database if relevant
+        if "sql" in vuln_type.lower() or "injection" in vuln_type.lower():
+            db_type_str = target_info.get("db_type", "mysql")
+            try:
+                db_type = DBType(db_type_str)
+                payloads = self.sqli_db.get_payloads(
+                    db_type=db_type, technique=SQLITechnique.UNION_BASED
+                )
+
+                exploit["payloads"] = [
+                    {"name": p.name, "payload": p.payload} for p in payloads[:3]
+                ]
+                exploit["severity"] = "critical"
+
+            except ValueError:
+                pass
+
+        # Use LLM for complex exploit development
+        if self.zen_orchestrator:
+            prompt = f"""
+            Develop a proof-of-concept exploit for:
+            Vulnerability: {vuln_type}
+            Target: {target_info}
+            
+            Provide:
+            1. Exploit technique overview
+            2. Step-by-step exploitation
+            3. Sample payload (safe for testing)
+            4. Detection methods
+            5. Mitigation recommendations
+            
+            FOR AUTHORIZED TESTING ONLY.
+            """
+
+            response = await self.zen_orchestrator.process(prompt)
+
+            exploit["techniques"] = response.content.split("\n")
+            exploit["summary"] = response.content[:200]
+
+        return exploit
+
+    async def generate_payload(self, payload_type: str, encoding: str = "raw") -> str:
+        """Generate a specific payload"""
+        if payload_type == "reverse_shell":
+            # Generate reverse shell payload
+            if self.zen_orchestrator:
+                prompt = """
+                Generate a bash reverse shell payload for testing.
+                Target: Linux
+                Safe for authorized testing only.
+                """
+                response = await self.zen_orchestrator.process(prompt)
+                return response.content
+
+        elif payload_type == "sqli":
+            # Get from SQL injection database
+            payloads = self.sqli_db.get_payloads(DBType.MYSQL)
+            if payloads:
+                return payloads[0].payload
+
+        return "# Payload generation failed"
+
+    async def collaborate_on_exploit(self, target_vuln: str, participants: List[str]):
+        """Collaborate with other agents on complex exploit development"""
+        # Request research from ResearchAgent
+        await self.send_message(
+            content=f"Need research on {target_vuln} for exploit development",
+            recipient="role:researcher",
+            msg_type="request_info",
+            requires_response=True,
+            context={"info_type": "exploit_vectors", "target": target_vuln},
+        )
+
+        # Request analysis from AnalysisAgent
+        await self.send_message(
+            content=f"Requesting attack path analysis for {target_vuln}",
+            recipient="role:analyst",
+            msg_type="analysis_request",
+            requires_response=True,
+            context={"analysis_type": "attack_path", "target": target_vuln},
+        )
+
+    async def execute_task(self, task: Dict) -> Dict:
+        """Execute exploit development task"""
+        task_type = task.get("type", "")
+        context = task.get("context", {})
+
+        logger.info(f"[ExploitAgent:{self.name}] Executing task: {task_type}")
+
+        if task_type == "exploit_development":
+            vuln = context.get("vulnerability", {})
+            exploit = await self._develop_exploit(
+                vuln.get("name", "unknown"), vuln.get("target_info", {})
+            )
+
+            # Share with coordinator
+            await self.send_message(
+                content="Exploit development complete",
+                recipient="role:coordinator",
+                msg_type="exploit_complete",
+                context={"exploit": exploit},
+            )
+
+            return {"status": "complete", "exploit": exploit, "agent": self.name}
+
+        return {"status": "unknown_task"}

agents/integration.py

@@ -0,0 +1,211 @@
+#!/usr/bin/env python3
+"""
+Integration Module - Connect Multi-Agent System to Zen AI Pentest
+Provides easy interface for pentesters to use agent collaboration
+Author: SHAdd0WTAka
+"""
+
+import asyncio
+import logging
+from typing import Any, Dict, List, Optional
+
+from .agent_orchestrator import AgentOrchestrator
+from .analysis_agent import AnalysisAgent
+from .exploit_agent import ExploitAgent
+from .research_agent import ResearchAgent
+
+logger = logging.getLogger("ZenAI.Agents")
+
+
+class AgentSystemIntegration:
+    """
+    High-level interface for using the multi-agent system
+    Pentesters can use this to coordinate agent research and analysis
+    """
+
+    def __init__(self, zen_orchestrator=None):
+        self.agent_orchestrator = AgentOrchestrator(zen_orchestrator)
+        self.zen_orchestrator = zen_orchestrator
+        self.initialized = False
+
+    async def initialize(self):
+        """Initialize the agent system with default agents"""
+        if self.initialized:
+            return
+
+        logger.info("[AgentIntegration] Initializing multi-agent system...")
+
+        # Create default agents
+        research_agent = ResearchAgent(
+            name="ResearchBot-Alpha",
+            orchestrator=self.agent_orchestrator,
+            zen_orchestrator=self.zen_orchestrator,
+        )
+
+        analysis_agent = AnalysisAgent(
+            name="AnalysisBot-Beta",
+            orchestrator=self.agent_orchestrator,
+            zen_orchestrator=self.zen_orchestrator,
+        )
+
+        exploit_agent = ExploitAgent(
+            name="ExploitBot-Gamma",
+            orchestrator=self.agent_orchestrator,
+            zen_orchestrator=self.zen_orchestrator,
+        )
+
+        # Register agents
+        self.agent_orchestrator.register_agent(research_agent)
+        self.agent_orchestrator.register_agent(analysis_agent)
+        self.agent_orchestrator.register_agent(exploit_agent)
+
+        # Start all agents
+        await self.agent_orchestrator.start_all()
+
+        self.initialized = True
+        logger.info("[AgentIntegration] Multi-agent system initialized with 3 agents")
+
+    async def conduct_research(self, topic: str, pentest_context: Dict) -> str:
+        """
+        Start coordinated research on a topic
+        Agents will work together and share findings
+        """
+        if not self.initialized:
+            await self.initialize()
+
+        thread_id = await self.agent_orchestrator.start_research_coordination(
+            topic=topic, pentest_context=pentest_context
+        )
+
+        logger.info(f"[AgentIntegration] Research started: {thread_id}")
+
+        # Wait a bit for agents to work
+        await asyncio.sleep(2)
+
+        return thread_id
+
+    async def analyze_target(self, target: str, findings: List[Dict]) -> Dict:
+        """
+        Have agents analyze a target collaboratively
+        """
+        if not self.initialized:
+            await self.initialize()
+
+        logger.info(f"[AgentIntegration] Starting collaborative analysis of {target}")
+
+        # Coordinate agents for vulnerability analysis
+        results = await self.agent_orchestrator.coordinate_agents(
+            task_type="vulnerability_analysis",
+            context={"target": target, "findings": findings},
+        )
+
+        return results
+
+    async def develop_exploits(self, vulnerabilities: List[Dict]) -> List[Dict]:
+        """
+        Have ExploitAgent develop exploits for vulnerabilities
+        """
+        if not self.initialized:
+            await self.initialize()
+
+        exploits = []
+
+        for vuln in vulnerabilities:
+            # Send exploit request to ExploitAgent
+            await self.agent_orchestrator.agents[
+                list(self.agent_orchestrator.agents.keys())[2]  # ExploitAgent
+            ].receive_message(
+                type(
+                    "Msg",
+                    (),
+                    {
+                        "sender": "user",
+                        "recipient": "ExploitBot-Gamma",
+                        "msg_type": "exploit_request",
+                        "context": {
+                            "vulnerability_type": vuln.get("name"),
+                            "target_info": vuln.get("target_info", {}),
+                        },
+                    },
+                )()
+            )
+
+        return exploits
+
+    async def facilitate_discussion(self, topic: str, rounds: int = 3) -> List[str]:
+        """
+        Have agents discuss a topic and share insights
+        Similar to Clawed/Moltbot conversations
+        """
+        if not self.initialized:
+            await self.initialize()
+
+        agent_ids = list(self.agent_orchestrator.agents.keys())
+
+        conversation = await self.agent_orchestrator.facilitate_conversation(
+            topic=topic, participants=agent_ids, rounds=rounds
+        )
+
+        # Extract conversation content
+        messages = [msg.content for msg in conversation]
+
+        return messages
+
+    def get_system_status(self) -> Dict:
+        """Get status of the multi-agent system"""
+        if not self.initialized:
+            return {"status": "not_initialized"}
+
+        return self.agent_orchestrator.get_system_status()
+
+    async def share_context(self, key: str, value: Any):
+        """Share context with all agents"""
+        if self.initialized:
+            await self.agent_orchestrator.update_shared_context(
+                key=key, value=value, source_agent_id="user"
+            )
+
+    async def shutdown(self):
+        """Shutdown the agent system gracefully"""
+        if self.initialized:
+            await self.agent_orchestrator.stop_all()
+            self.initialized = False
+            logger.info("[AgentIntegration] Agent system shutdown complete")
+
+
+# Convenience functions for direct use
+
+
+async def start_collaborative_research(
+    topic: str, zen_orchestrator=None
+) -> AgentSystemIntegration:
+    """
+    Quick-start function to begin collaborative research
+
+    Usage:
+        integration = await start_collaborative_research("WordPress vulnerabilities")
+        # Let agents work...
+        status = integration.get_system_status()
+    """
+    integration = AgentSystemIntegration(zen_orchestrator)
+    await integration.initialize()
+    await integration.conduct_research(
+        topic=topic, pentest_context={"target_type": "web", "scope": "full"}
+    )
+    return integration
+
+
+async def multi_agent_analysis(
+    target: str, findings: List[Dict], zen_orchestrator=None
+) -> Dict:
+    """
+    One-shot function for multi-agent analysis
+
+    Usage:
+        results = await multi_agent_analysis("example.com", nmap_findings)
+        print(results["shared_findings"])
+    """
+    integration = AgentSystemIntegration(zen_orchestrator)
+    results = await integration.analyze_target(target, findings)
+    await integration.shutdown()
+    return results