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

agents/__init__.py

@@ -0,0 +1,28 @@
+"""
+Multi-Agent Collaboration System for Zen AI Pentest
+Inspired by Clawed/Moltbot architecture
+Agents can communicate, share context, and coordinate research
+Author: SHAdd0WTAka
+"""
+
+from .agent_base import AgentMessage, AgentRole, BaseAgent
+from .agent_orchestrator import AgentOrchestrator
+from .analysis_agent import AnalysisAgent
+from .exploit_agent import ExploitAgent
+from .post_scan_agent import (PentestLoot, PostScanAgent, VerifiedFinding,
+                              run_post_scan_workflow)
+from .research_agent import ResearchAgent
+
+__all__ = [
+    "BaseAgent",
+    "AgentRole",
+    "AgentMessage",
+    "AgentOrchestrator",
+    "ResearchAgent",
+    "ExploitAgent",
+    "AnalysisAgent",
+    "PostScanAgent",
+    "run_post_scan_workflow",
+    "VerifiedFinding",
+    "PentestLoot",
+]

agents/agent_base.py

@@ -0,0 +1,239 @@
+#!/usr/bin/env python3
+"""
+Base Agent Class for Multi-Agent System
+Provides messaging, context sharing, and role-based capabilities
+Author: SHAdd0WTAka
+"""
+
+import asyncio
+import logging
+import uuid
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from datetime import datetime
+from enum import Enum
+from typing import Any, Callable, Dict, List, Optional
+
+logger = logging.getLogger("ZenAI.Agents")
+
+
+class AgentRole(Enum):
+    """Agent roles for specialization"""
+
+    RESEARCHER = "researcher"  # Gathers information, reconnaissance
+    ANALYST = "analyst"  # Analyzes data, finds patterns
+    EXPLOIT = "exploit"  # Develops exploits, payloads
+    COORDINATOR = "coordinator"  # Manages workflow between agents
+    REPORTER = "reporter"  # Generates reports, summaries
+    POST_EXPLOITATION = (
+        "post_exploit"  # Post-scan workflow: verification, evidence, cleanup
+    )
+
+
+class AgentState(Enum):
+    """Operational states for agents"""
+
+    IDLE = "idle"
+    BUSY = "busy"
+    WAITING = "waiting"
+    COMPLETED = "completed"
+    FAILED = "failed"
+    STOPPED = "stopped"
+
+
+
+@dataclass
+class AgentMessage:
+    """Message format for inter-agent communication"""
+
+    id: str = field(default_factory=lambda: str(uuid.uuid4())[:8])
+    sender: str = ""
+    recipient: str = ""  # "all" for broadcast, specific ID for direct
+    msg_type: str = "chat"  # chat, task, result, request, response
+    content: str = ""
+    context: Dict[str, Any] = field(default_factory=dict)
+    timestamp: str = field(default_factory=lambda: datetime.now().isoformat())
+    priority: int = 1  # 1=low, 2=medium, 3=high, 4=critical
+    requires_response: bool = False
+
+    def to_dict(self) -> Dict:
+        return {
+            "id": self.id,
+            "sender": self.sender,
+            "recipient": self.recipient,
+            "type": self.msg_type,
+            "content": self.content,
+            "context": self.context,
+            "timestamp": self.timestamp,
+            "priority": self.priority,
+            "requires_response": self.requires_response,
+        }
+
+
+class BaseAgent(ABC):
+    """
+    Base class for all agents in the system
+    Provides messaging, context management, and coordination capabilities
+    """
+
+    def __init__(self, name: str, role: AgentRole, orchestrator=None):
+        self.id = str(uuid.uuid4())[:8]
+        self.name = name
+        self.role = role
+        self.orchestrator = orchestrator
+        self.message_queue: asyncio.Queue = asyncio.Queue()
+        self.inbox: List[AgentMessage] = []
+        self.context: Dict[str, Any] = {}  # Shared workspace
+        self.memory: List[Dict] = []  # Agent's memory of interactions
+        self.handlers: Dict[str, Callable] = {}
+        self.running = False
+        self.task: Optional[asyncio.Task] = None
+
+        logger.info(f"[Agent] Initialized {self.name} ({self.role.value}) [{self.id}]")
+
+    def register_handler(self, msg_type: str, handler: Callable):
+        """Register a handler for specific message types"""
+        self.handlers[msg_type] = handler
+
+    async def send_message(
+        self,
+        content: str,
+        recipient: str = "all",
+        msg_type: str = "chat",
+        priority: int = 1,
+        requires_response: bool = False,
+        context: Dict = None,
+    ) -> str:
+        """
+        Send a message to another agent or broadcast
+        Returns message ID for tracking
+        """
+        msg = AgentMessage(
+            sender=f"{self.name}[{self.id}]",
+            recipient=recipient,
+            msg_type=msg_type,
+            content=content,
+            priority=priority,
+            requires_response=requires_response,
+            context=context or {},
+        )
+
+        if self.orchestrator:
+            await self.orchestrator.route_message(msg)
+            logger.debug(f"[Agent:{self.name}] Sent {msg_type} to {recipient}")
+        else:
+            # Direct send if no orchestrator
+            await self.message_queue.put(msg)
+
+        return msg.id
+
+    async def receive_message(self, msg: AgentMessage):
+        """Receive a message into the queue"""
+        await self.message_queue.put(msg)
+        self.inbox.append(msg)
+
+        # Store in memory
+        self.memory.append(
+            {"type": "received", "message": msg.to_dict(), "processed": False}
+        )
+
+        logger.debug(f"[Agent:{self.name}] Received message from {msg.sender}")
+
+    async def process_messages(self):
+        """Main message processing loop"""
+        self.running = True
+
+        while self.running:
+            try:
+                # Get message with timeout to allow checking running flag
+                msg = await asyncio.wait_for(self.message_queue.get(), timeout=1.0)
+
+                # Process based on type
+                if msg.msg_type in self.handlers:
+                    await self.handlers[msg.msg_type](msg)
+                else:
+                    await self.handle_message(msg)
+
+                # Mark as processed in memory
+                for mem in self.memory:
+                    if mem["type"] == "received" and mem["message"]["id"] == msg.id:
+                        mem["processed"] = True
+
+            except asyncio.TimeoutError:
+                continue
+            except Exception as e:
+                logger.error(f"[Agent:{self.name}] Error processing message: {e}")
+
+    async def handle_message(self, msg: AgentMessage):
+        """Default message handler - override in subclasses"""
+        logger.info(f"[Agent:{self.name}] Handling {msg.msg_type} from {msg.sender}")
+
+        # Default chat response
+        if msg.msg_type == "chat":
+            response = f"Acknowledged: {msg.content[:50]}..."
+            if msg.requires_response:
+                await self.send_message(
+                    content=response, recipient=msg.sender, msg_type="response"
+                )
+
+    def update_context(self, key: str, value: Any, share: bool = False):
+        """
+        Update agent's local context
+        If share=True, broadcast to all agents via orchestrator
+        """
+        self.context[key] = value
+
+        if share and self.orchestrator:
+            asyncio.create_task(
+                self.orchestrator.update_shared_context(key, value, self.id)
+            )
+
+    def get_context(self, key: str) -> Any:
+        """Get value from context"""
+        return self.context.get(key)
+
+    def share_findings(self, findings: Dict):
+        """Share findings with all agents"""
+        if self.orchestrator:
+            asyncio.create_task(
+                self.send_message(
+                    content=f"Sharing {len(findings)} findings",
+                    msg_type="findings",
+                    context={"findings": findings},
+                    priority=2,
+                )
+            )
+
+    @abstractmethod
+    async def execute_task(self, task: Dict) -> Dict:
+        """Execute a specific task - must be implemented by subclasses"""
+        pass
+
+    async def start(self):
+        """Start the agent's message processing loop"""
+        self.task = asyncio.create_task(self.process_messages())
+        logger.info(f"[Agent:{self.name}] Started")
+
+    async def stop(self):
+        """Stop the agent gracefully"""
+        self.running = False
+        if self.task:
+            self.task.cancel()
+            try:
+                await self.task
+            except asyncio.CancelledError:
+                pass
+        logger.info(f"[Agent:{self.name}] Stopped")
+
+    def get_status(self) -> Dict:
+        """Get agent status"""
+        return {
+            "id": self.id,
+            "name": self.name,
+            "role": self.role.value,
+            "running": self.running,
+            "queue_size": self.message_queue.qsize(),
+            "inbox_count": len(self.inbox),
+            "memory_entries": len(self.memory),
+            "context_keys": list(self.context.keys()),
+        }

agents/agent_orchestrator.py

@@ -0,0 +1,346 @@
+#!/usr/bin/env python3
+"""
+Agent Orchestrator
+Manages multiple agents, routes messages, coordinates research
+Like Clawed/Moltbot but for penetration testing
+Author: SHAdd0WTAka
+"""
+
+import asyncio
+import logging
+from datetime import datetime
+from typing import Any, Dict, List, Optional
+
+from .agent_base import AgentMessage, AgentRole, BaseAgent
+
+logger = logging.getLogger("ZenAI.Agents")
+
+
+class AgentOrchestrator:
+    """
+    Central coordinator for multi-agent system
+    Manages agent lifecycle, message routing, and context sharing
+    """
+
+    def __init__(self, zen_orchestrator=None):
+        self.agents: Dict[str, BaseAgent] = {}
+        self.agent_by_role: Dict[AgentRole, List[BaseAgent]] = {}
+        self.shared_context: Dict[str, Any] = {}
+        self.message_history: List[AgentMessage] = []
+        self.conversation_threads: Dict[str, List[str]] = {}  # thread_id -> message_ids
+        self.zen_orchestrator = zen_orchestrator  # Link to LLM orchestrator
+        self.running = False
+        self.research_coordination: Dict[str, Any] = {}
+
+    def register_agent(self, agent: BaseAgent):
+        """Register an agent with the orchestrator"""
+        self.agents[agent.id] = agent
+        agent.orchestrator = self
+
+        # Index by role
+        if agent.role not in self.agent_by_role:
+            self.agent_by_role[agent.role] = []
+        self.agent_by_role[agent.role].append(agent)
+
+        logger.info(
+            f"[Orchestrator] Registered agent {agent.name} ({agent.role.value})"
+        )
+
+    def unregister_agent(self, agent_id: str):
+        """Remove an agent from the system"""
+        if agent_id in self.agents:
+            agent = self.agents[agent_id]
+
+            # Remove from role index
+            if agent.role in self.agent_by_role:
+                self.agent_by_role[agent.role] = [
+                    a for a in self.agent_by_role[agent.role] if a.id != agent_id
+                ]
+
+            del self.agents[agent_id]
+            logger.info(f"[Orchestrator] Unregistered agent {agent_id}")
+
+    async def route_message(self, msg: AgentMessage):
+        """Route a message to the appropriate recipient(s)"""
+        self.message_history.append(msg)
+
+        if msg.recipient == "all":
+            # Broadcast to all agents except sender
+            for agent_id, agent in self.agents.items():
+                if f"[{agent_id}]" not in msg.sender:
+                    await agent.receive_message(msg)
+
+        elif msg.recipient.startswith("role:"):
+            # Send to all agents of a specific role
+            role_str = msg.recipient.split(":")[1]
+            try:
+                role = AgentRole(role_str)
+                for agent in self.agent_by_role.get(role, []):
+                    if f"[{agent.id}]" not in msg.sender:
+                        await agent.receive_message(msg)
+            except ValueError:
+                logger.error(f"[Orchestrator] Invalid role: {role_str}")
+
+        else:
+            # Direct message to specific agent
+            # Parse agent ID from recipient string like "AgentName[ID]"
+            if "[" in msg.recipient and "]" in msg.recipient:
+                agent_id = msg.recipient.split("[")[1].split("]")[0]
+                if agent_id in self.agents:
+                    await self.agents[agent_id].receive_message(msg)
+
+    async def update_shared_context(self, key: str, value: Any, source_agent_id: str):
+        """Update shared context and notify all agents"""
+        self.shared_context[key] = {
+            "value": value,
+            "updated_by": source_agent_id,
+            "timestamp": datetime.now().isoformat(),
+        }
+
+        # Notify all agents of context update
+        for agent_id, agent in self.agents.items():
+            if agent_id != source_agent_id:
+                await agent.receive_message(
+                    AgentMessage(
+                        sender="orchestrator",
+                        recipient=f"{agent.name}[{agent.id}]",
+                        msg_type="context_update",
+                        content=f"Context updated: {key}",
+                        context={"key": key, "value": value},
+                    )
+                )
+
+    def get_shared_context(self, key: str = None) -> Any:
+        """Get value from shared context"""
+        if key:
+            return self.shared_context.get(key, {}).get("value")
+        return self.shared_context
+
+    async def start_research_coordination(
+        self, topic: str, pentest_context: Dict
+    ) -> str:
+        """
+        Coordinate multi-agent research on a topic
+        Returns thread ID for tracking
+        """
+        thread_id = f"research_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
+
+        self.research_coordination[thread_id] = {
+            "topic": topic,
+            "context": pentest_context,
+            "started": datetime.now().isoformat(),
+            "agents_involved": [],
+            "findings": [],
+            "status": "active",
+        }
+
+        # Notify all research agents
+        for agent in self.agent_by_role.get(AgentRole.RESEARCHER, []):
+            await agent.receive_message(
+                AgentMessage(
+                    sender="orchestrator",
+                    recipient=f"{agent.name}[{agent.id}]",
+                    msg_type="research_task",
+                    content=f"Research coordination started: {topic}",
+                    context={
+                        "thread_id": thread_id,
+                        "pentest_context": pentest_context,
+                        "shared_context": self.shared_context,
+                    },
+                    priority=3,
+                )
+            )
+            self.research_coordination[thread_id]["agents_involved"].append(agent.id)
+
+        logger.info(f"[Orchestrator] Started research coordination: {thread_id}")
+        return thread_id
+
+    async def coordinate_agents(self, task_type: str, context: Dict) -> Dict:
+        """
+        Coordinate multiple agents for a complex task
+        Agents will communicate and share findings
+        """
+        results = {
+            "task_type": task_type,
+            "started": datetime.now().isoformat(),
+            "agent_responses": {},
+        }
+
+        # Determine which agents to involve based on task
+        involved_roles = []
+
+        if task_type == "reconnaissance":
+            involved_roles = [AgentRole.RESEARCHER, AgentRole.ANALYST]
+        elif task_type == "vulnerability_analysis":
+            involved_roles = [AgentRole.ANALYST, AgentRole.EXPLOIT]
+        elif task_type == "exploit_development":
+            involved_roles = [AgentRole.EXPLOIT, AgentRole.ANALYST]
+        elif task_type == "full_assessment":
+            involved_roles = [
+                AgentRole.RESEARCHER,
+                AgentRole.ANALYST,
+                AgentRole.EXPLOIT,
+            ]
+
+        # Start all involved agents
+        tasks = []
+        for role in involved_roles:
+            for agent in self.agent_by_role.get(role, []):
+                agent_task = asyncio.create_task(
+                    agent.execute_task(
+                        {
+                            "type": task_type,
+                            "context": context,
+                            "shared_context": self.shared_context,
+                        }
+                    )
+                )
+                tasks.append((agent.id, agent_task))
+
+        # Wait for all agents to complete
+        for agent_id, task in tasks:
+            try:
+                result = await asyncio.wait_for(task, timeout=300)
+                results["agent_responses"][agent_id] = result
+            except asyncio.TimeoutError:
+                results["agent_responses"][agent_id] = {"error": "timeout"}
+
+        # Agents should have communicated with each other during execution
+        # Collect their findings from shared context
+        results["shared_findings"] = self.shared_context
+        results["completed"] = datetime.now().isoformat()
+
+        return results
+
+    async def facilitate_conversation(
+        self, topic: str, participants: List[str], rounds: int = 3
+    ) -> List[AgentMessage]:
+        """
+        Facilitate a multi-round conversation between agents
+        Similar to Clawed/Moltbot group discussions
+        """
+        conversation = []
+
+        for round_num in range(rounds):
+            logger.info(f"[Orchestrator] Conversation round {round_num + 1}/{rounds}")
+
+            for participant_id in participants:
+                if participant_id in self.agents:
+                    agent = self.agents[participant_id]
+
+                    # Request agent's thoughts on topic
+                    msg = AgentMessage(
+                        sender="orchestrator",
+                        recipient=f"{agent.name}[{agent.id}]",
+                        msg_type="conversation_round",
+                        content=f"Round {round_num + 1}: Share your thoughts on {topic}",
+                        context={
+                            "round": round_num + 1,
+                            "topic": topic,
+                            "conversation_so_far": [m.to_dict() for m in conversation],
+                            "shared_context": self.shared_context,
+                        },
+                        requires_response=True,
+                    )
+
+                    await agent.receive_message(msg)
+
+        return conversation
+
+    async def start_all(self):
+        """Start all registered agents"""
+        self.running = True
+        for agent in self.agents.values():
+            await agent.start()
+
+    async def stop_all(self):
+        """Stop all agents gracefully"""
+        self.running = False
+        for agent in self.agents.values():
+            await agent.stop()
+
+    async def execute_post_scan_workflow(
+        self, target: str, scan_results: Dict[str, Any]
+    ) -> Dict[str, Any]:
+        """
+        Execute the complete post-scan pentester workflow
+        This runs automatically after every scan to ensure professional standards
+
+        Phases:
+        1. Manual Verification (false positive elimination)
+        2. Vulnerability Validation
+        3. Exploitation Attempts
+        4. Post-Exploitation (privilege escalation, lateral movement)
+        5. Evidence Collection
+        6. Loot Documentation
+        7. Cleanup & Restoration
+        8. Report Preparation
+        """
+        from .post_scan_agent import PostScanAgent
+
+        logger.info(f"[Orchestrator] Starting post-scan workflow for {target}")
+        print(f"\n[Post-Scan Workflow] Initiating professional pentest follow-up...")
+
+        # Create and run post-scan agent
+        post_scan_agent = PostScanAgent()
+
+        # Extract findings from scan results
+        findings = scan_results.get("findings", [])
+        if not findings:
+            # Create sample findings if none exist
+            findings = self._generate_sample_findings(target)
+
+        # Execute the workflow
+        results = await post_scan_agent.run(target, findings)
+
+        # Store in shared context
+        await self.update_shared_context(f"post_scan_{target}", results, "orchestrator")
+
+        logger.info(f"[Orchestrator] Post-scan workflow complete for {target}")
+        return results
+
+    def _generate_sample_findings(self, target: str) -> List[Dict]:
+        """Generate sample findings for demonstration"""
+        return [
+            {
+                "id": "CVE-2021-44228",
+                "title": "Log4j Remote Code Execution",
+                "severity": "critical",
+                "cvss_score": 10.0,
+                "description": "Log4Shell vulnerability allows RCE",
+                "port": 8080,
+                "service": "http",
+            },
+            {
+                "id": "WEAK_SSH",
+                "title": "SSH Weak Cipher Suites",
+                "severity": "medium",
+                "cvss_score": 5.3,
+                "description": "SSH supports weak ciphers",
+                "port": 22,
+                "service": "ssh",
+            },
+            {
+                "id": "DEFAULT_CREDS",
+                "title": "Default Credentials Detected",
+                "severity": "high",
+                "cvss_score": 8.1,
+                "description": "Default admin/admin credentials work",
+                "port": 80,
+                "service": "http",
+            },
+        ]
+
+    def get_system_status(self) -> Dict:
+        """Get status of entire multi-agent system"""
+        return {
+            "agents": {
+                agent_id: agent.get_status() for agent_id, agent in self.agents.items()
+            },
+            "shared_context_keys": list(self.shared_context.keys()),
+            "message_count": len(self.message_history),
+            "active_research": list(self.research_coordination.keys()),
+            "role_distribution": {
+                role.value: len(agents) for role, agents in self.agent_by_role.items()
+            },
+        }