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

agents/react_agent_vm.py

@@ -0,0 +1,298 @@
+"""
+ReAct Agent mit VirtualBox VM-Integration
+
+Erweiterung des ReAct Agents für isolierte VM-basierte Pentests.
+"""
+
+from typing import List, TypedDict, Annotated, Literal
+import logging
+
+from langchain_core.tools import tool
+from langchain_core.messages import BaseMessage, HumanMessage, AIMessage, ToolMessage, SystemMessage
+from langchain_core.prompts import ChatPromptTemplate
+from langgraph.graph import StateGraph, START, END, add_messages
+from langgraph.checkpoint.memory import MemorySaver
+
+from .react_agent import ReActAgent, ReActAgentConfig, AgentState
+from ..virtualization.vm_manager import VirtualBoxManager, PentestSandbox
+
+
+logger = logging.getLogger(__name__)
+
+
+class VMAgentConfig(ReActAgentConfig):
+    """Erweiterte Konfiguration für VM-basierten Agent"""
+    use_vm: bool = True
+    vm_name: str = "kali-pentest"
+    vm_username: str = "kali"
+    vm_password: str = "kali"
+    auto_restore_snapshot: bool = True
+
+
+class VMReActAgent(ReActAgent):
+    """
+    ReAct Agent mit VM-Backend.
+    
+    Führt alle Tools in einer isolierten Kali Linux VM aus
+    für maximale Sicherheit und Reproduzierbarkeit.
+    """
+    
+    def __init__(self, config: VMAgentConfig = None):
+        self.vm_config = config or VMAgentConfig()
+        
+        # VM Manager initialisieren
+        if self.vm_config.use_vm:
+            self.vbox = VirtualBoxManager()
+            self.sandbox = PentestSandbox(self.vbox)
+            self.session_id = None
+        
+        # Parent init ohne Tools (wir überschreiben sie)
+        self.config = self.vm_config
+        self.llm = None  # Wird in _build_graph gesetzt
+        self.tools = []
+        self.graph = self._build_graph_with_vm()
+        
+        logger.info(f"VMReActAgent initialisiert (VM: {self.vm_config.vm_name})")
+    
+    def _build_graph_with_vm(self):
+        """Baut Graph mit VM-Tools"""
+        from ..core.llm_backend import LLMBackend
+        
+        self.llm = LLMBackend(model=self.config.llm_model)
+        
+        # VM-basierte Tools erstellen
+        self.tools = self._create_vm_tools()
+        llm_with_tools = self.llm.bind_tools(self.tools)
+        
+        system_prompt = """Du bist ein autonomer Pentest-Agent mit VM-Integration.
+
+Deine Aufgabe:
+1. Führe alle Tests in einer isolierten VM durch
+2. Nutze Snapshots für Clean-State-Workflows
+3. Dokumentiere alle Findings strukturiert
+4. Fasse Ergebnisse am Ende zusammen
+
+SICHERHEIT:
+- Alle Scans laufen in der VM 'kali-pentest'
+- Snapshots werden automatisch verwaltet
+- Keine Auswirkungen auf das Host-System
+
+Wenn fertig, gib eine finale Zusammenfassung."""
+
+        def agent_node(state: AgentState) -> AgentState:
+            if state["iteration"] >= state["max_iterations"]:
+                return {
+                    **state,
+                    "messages": state["messages"] + [
+                        AIMessage(content="Maximale Iterationen erreicht.")
+                    ],
+                    "status": "completed"
+                }
+            
+            messages = [SystemMessage(content=system_prompt)] + state["messages"]
+            response = llm_with_tools.invoke(messages)
+            
+            return {
+                **state,
+                "messages": [response],
+                "iteration": state["iteration"] + 1
+            }
+        
+        def tools_node(state: AgentState) -> AgentState:
+            last_message = state["messages"][-1]
+            
+            if not hasattr(last_message, 'tool_calls') or not last_message.tool_calls:
+                return state
+            
+            tool_messages = []
+            new_findings = []
+            
+            for tool_call in last_message.tool_calls:
+                tool_name = tool_call["name"]
+                args = tool_call["args"]
+                tool_id = tool_call["id"]
+                
+                logger.info(f"VM-Tool-Aufruf: {tool_name}")
+                
+                try:
+                    tool_func = {t.name: t for t in self.tools}.get(tool_name)
+                    if tool_func:
+                        result = tool_func.invoke(args)
+                        tool_messages.append(ToolMessage(content=result, tool_call_id=tool_id))
+                        new_findings.append({
+                            "tool": tool_name,
+                            "args": args,
+                            "result": result
+                        })
+                    else:
+                        tool_messages.append(ToolMessage(
+                            content=f"Tool {tool_name} nicht gefunden", 
+                            tool_call_id=tool_id
+                        ))
+                except Exception as e:
+                    error_msg = f"VM-Tool-Fehler: {str(e)}"
+                    logger.error(error_msg)
+                    tool_messages.append(ToolMessage(content=error_msg, tool_call_id=tool_id))
+            
+            return {
+                **state,
+                "messages": tool_messages,
+                "findings": state["findings"] + new_findings
+            }
+        
+        def should_continue(state: AgentState):
+            last_message = state["messages"][-1]
+            
+            if state.get("status") == "completed":
+                return "end"
+            
+            if hasattr(last_message, 'tool_calls') and last_message.tool_calls:
+                return "tools"
+            
+            if isinstance(last_message, ToolMessage):
+                return "agent"
+            
+            return "end"
+        
+        workflow = StateGraph(AgentState)
+        workflow.add_node("agent", agent_node)
+        workflow.add_node("tools", tools_node)
+        workflow.add_edge(START, "agent")
+        workflow.add_conditional_edges("agent", should_continue, 
+                                      {"tools": "tools", "agent": "agent", "end": END})
+        workflow.add_conditional_edges("tools", should_continue,
+                                      {"tools": "tools", "agent": "agent", "end": END})
+        
+        checkpointer = MemorySaver()
+        return workflow.compile(checkpointer=checkpointer)
+    
+    def _create_vm_tools(self) -> List:
+        """Erstellt VM-basierte Tools"""
+        from langchain_core.tools import tool
+        
+        config = self.vm_config
+        
+        @tool
+        def vm_nmap_scan(target: str, ports: str = "top-100") -> str:
+            """Scannt Ports in VM mit Nmap"""
+            if not self.session_id:
+                return "Keine aktive Session"
+            
+            port_arg = "--top-ports 100" if ports == "top-100" else f"-p {ports}"
+            exit_code, stdout, stderr = self.sandbox.execute_tool(
+                self.session_id, "nmap", f"{port_arg} {target}",
+                config.vm_username, config.vm_password
+            )
+            return stdout if exit_code == 0 else stderr
+        
+        @tool
+        def vm_nuclei_scan(target: str, severity: str = "critical,high") -> str:
+            """Scannt nach CVEs in VM mit Nuclei"""
+            if not self.session_id:
+                return "Keine aktive Session"
+            
+            exit_code, stdout, stderr = self.sandbox.execute_tool(
+                self.session_id, "nuclei",
+                f"-u {target} -s {severity}",
+                config.vm_username, config.vm_password
+            )
+            return stdout if exit_code == 0 else stderr
+        
+        @tool
+        def vm_gobuster_scan(target: str, wordlist: str = "common.txt") -> str:
+            """Directory Enumeration in VM mit Gobuster"""
+            if not self.session_id:
+                return "Keine aktive Session"
+            
+            exit_code, stdout, stderr = self.sandbox.execute_tool(
+                self.session_id, "gobuster",
+                f"dir -u {target} -w /usr/share/wordlists/dirb/{wordlist}",
+                config.vm_username, config.vm_password
+            )
+            return stdout if exit_code == 0 else stderr
+        
+        @tool
+        def restore_vm_snapshot(snapshot_name: str = "clean_state") -> str:
+            """Stellt VM-Snapshot wieder her"""
+            success = self.vbox.restore_snapshot(config.vm_name, snapshot_name)
+            return f"Snapshot {snapshot_name} wiederhergestellt" if success else "Fehler"
+        
+        return [vm_nmap_scan, vm_nuclei_scan, vm_gobuster_scan, restore_vm_snapshot]
+    
+    def run(self, target: str, objective: str = "comprehensive scan") -> dict:
+        """Führt VM-basierten Pentest durch"""
+        
+        # VM Session starten
+        if self.vm_config.use_vm:
+            import uuid
+            self.session_id = f"vm_pentest_{uuid.uuid4().hex[:8]}"
+            
+            logger.info(f"Starte VM-Session: {self.session_id}")
+            
+            if self.vm_config.auto_restore_snapshot:
+                logger.info("Stelle clean_state Snapshot wieder her...")
+                self.vbox.restore_snapshot(self.vm_config.vm_name, "clean_state")
+            
+            if not self.sandbox.create_session(self.session_id, self.vm_config.vm_name):
+                raise RuntimeError("Konnte VM-Session nicht starten")
+        
+        try:
+            # Parent run aufrufen
+            initial_state = {
+                "messages": [HumanMessage(content=f"{objective} on {target}")],
+                "findings": [],
+                "target": target,
+                "iteration": 0,
+                "max_iterations": self.config.max_iterations,
+                "status": "running"
+            }
+            
+            result = self.graph.invoke(
+                initial_state,
+                config={"configurable": {"thread_id": self.session_id}}
+            )
+            
+            return {
+                "findings": result["findings"],
+                "final_message": result["messages"][-1].content if result["messages"] else "",
+                "iterations": result["iteration"],
+                "status": result["status"],
+                "target": target,
+                "session_id": self.session_id
+            }
+            
+        finally:
+            # Aufräumen
+            if self.session_id and self.vm_config.use_vm:
+                logger.info("Beende VM-Session...")
+                self.sandbox.end_session(self.session_id)
+                self.session_id = None
+
+
+# Factory
+def get_vm_agent(config: VMAgentConfig = None) -> VMReActAgent:
+    """Gibt VM-basierten Agent zurück"""
+    return VMReActAgent(config)
+
+
+if __name__ == "__main__":
+    logging.basicConfig(level=logging.INFO)
+    
+    config = VMAgentConfig(
+        use_vm=True,
+        vm_name="kali-pentest",
+        max_iterations=5,
+        auto_restore_snapshot=True
+    )
+    
+    agent = VMReActAgent(config)
+    result = agent.run("scanme.nmap.org", "Port scan and vulnerability assessment")
+    
+    print("\n" + "=" * 70)
+    print("VM-BASIERTER PENTEST ABGESCHLOSSEN")
+    print("=" * 70)
+    print(f"Session: {result.get('session_id')}")
+    print(f"Iterations: {result['iterations']}")
+    print(f"Findings: {len(result['findings'])}")
+    print("\nFinal Analysis:")
+    print(result['final_message'][:500])

agents/research_agent.py

@@ -0,0 +1,176 @@
+#!/usr/bin/env python3
+"""
+Research Agent - Specializes in reconnaissance and information gathering
+Part of the Multi-Agent Collaboration System
+Author: SHAdd0WTAka
+"""
+
+import asyncio
+import logging
+from typing import Any, Dict, List
+
+from modules.cve_database import CVEDatabase
+from modules.sql_injection_db import SQLInjectionDatabase
+
+from .agent_base import AgentMessage, AgentRole, BaseAgent
+
+logger = logging.getLogger("ZenAI.Agents")
+
+
+class ResearchAgent(BaseAgent):
+    """
+    Research Agent - Gathers information, performs reconnaissance
+    Can research CVEs, exploits, techniques, and targets
+    """
+
+    def __init__(self, name: str, orchestrator=None, zen_orchestrator=None):
+        super().__init__(name, AgentRole.RESEARCHER, orchestrator)
+        self.zen_orchestrator = zen_orchestrator  # LLM orchestrator
+        self.cve_db = CVEDatabase()
+        self.sqli_db = SQLInjectionDatabase()
+        self.current_research = {}
+
+        # Register message handlers
+        self.register_handler("research_task", self._handle_research_task)
+        self.register_handler("request_info", self._handle_info_request)
+        self.register_handler("findings", self._handle_findings)
+
+    async def _handle_research_task(self, msg: AgentMessage):
+        """Handle incoming research tasks"""
+        logger.info(f"[ResearchAgent:{self.name}] Received research task")
+
+        task_context = msg.context.get("pentest_context", {})
+        thread_id = msg.context.get("thread_id", "unknown")
+
+        # Acknowledge receipt
+        await self.send_message(
+            content=f"Starting research on {task_context.get('target', 'unknown target')}",
+            recipient=msg.sender,
+            msg_type="research_status",
+            context={"thread_id": thread_id, "status": "started"},
+        )
+
+        # Perform research
+        findings = await self._perform_research(task_context)
+
+        # Share findings with all agents
+        await self.send_message(
+            content=f"Research complete. Found {len(findings)} relevant items.",
+            recipient="all",
+            msg_type="findings",
+            priority=2,
+            context={
+                "thread_id": thread_id,
+                "findings": findings,
+                "researcher": self.name,
+            },
+        )
+
+        # Update shared context
+        self.update_context(f"research_{thread_id}", findings, share=True)
+
+    async def _handle_info_request(self, msg: AgentMessage):
+        """Handle requests for specific information"""
+        request_type = msg.context.get("info_type")
+        query = msg.content
+
+        logger.info(f"[ResearchAgent:{self.name}] Info request: {request_type}")
+
+        result = None
+
+        if request_type == "cve":
+            result = self.cve_db.search_cve(query)
+        elif request_type == "ransomware":
+            result = self.cve_db.search_ransomware(query)
+        elif request_type == "sqli_payload":
+            db_type = msg.context.get("db_type")
+            result = self.sqli_db.get_payloads(db_type=db_type)
+
+        if result:
+            await self.send_message(
+                content=f"Found information about {query}",
+                recipient=msg.sender,
+                msg_type="response",
+                context={"result": result, "query": query},
+            )
+
+    async def _handle_findings(self, msg: AgentMessage):
+        """Handle findings from other agents - cross-reference with research"""
+        findings = msg.context.get("findings", {})
+
+        # If findings contain CVEs, enrich them with database info
+        if "cves" in findings:
+            enriched = []
+            for cve_id in findings["cves"]:
+                cve_info = self.cve_db.search_cve(cve_id)
+                if cve_info:
+                    enriched.append(cve_info)
+
+            if enriched:
+                await self.send_message(
+                    content=f"Enriched {len(enriched)} CVEs with database information",
+                    recipient=msg.sender,
+                    msg_type="enrichment",
+                    context={"enriched_cves": enriched},
+                )
+
+    async def _perform_research(self, context: Dict) -> List[Dict]:
+        """Perform actual research based on context"""
+        findings = []
+        target = context.get("target", "")
+
+        # Research 1: Check for known CVEs related to target tech
+        technologies = context.get("technologies", [])
+        for tech in technologies:
+            # Use LLM to find relevant CVEs
+            if self.zen_orchestrator:
+                prompt = f"""
+                List known CVEs for {tech} that are commonly exploited by ransomware.
+                Return CVE IDs only, one per line.
+                """
+                response = await self.zen_orchestrator.process(prompt)
+
+                # Parse CVEs from response
+                import re
+
+                cves = re.findall(r"CVE-\d{4}-\d{4,}", response.content)
+
+                for cve_id in cves[:5]:  # Limit to top 5
+                    cve_data = self.cve_db.search_cve(cve_id)
+                    if cve_data:
+                        findings.append(
+                            {"type": "cve", "data": cve_data, "source": "llm_research"}
+                        )
+
+        # Research 2: Check for ransomware using these technologies
+        ransomware = self.cve_db.list_all_ransomware()
+        for rw in ransomware[:3]:  # Check top 3
+            findings.append({"type": "ransomware", "data": rw, "source": "database"})
+
+        return findings
+
+    async def execute_task(self, task: Dict) -> Dict:
+        """Execute a research task"""
+        task_type = task.get("type", "")
+        context = task.get("context", {})
+
+        logger.info(f"[ResearchAgent:{self.name}] Executing task: {task_type}")
+
+        if task_type == "reconnaissance":
+            findings = await self._perform_research(context)
+
+            # Share with other agents
+            await self.send_message(
+                content="Reconnaissance findings ready",
+                recipient="role:analyst",
+                msg_type="findings",
+                context={"findings": findings},
+            )
+
+            return {
+                "status": "complete",
+                "findings_count": len(findings),
+                "agent": self.name,
+            }
+
+        return {"status": "unknown_task"}

api/__init__.py

@@ -0,0 +1,11 @@
+"""
+Zen AI Pentest - REST API Module
+
+A production-ready API for penetration testing operations.
+Built with FastAPI for high performance and automatic documentation.
+
+Author: SHAdd0WTAka + Kimi AI
+"""
+
+__version__ = "1.0.0"
+__all__ = ["app", "create_app"]

api/auth.py

@@ -0,0 +1,123 @@
+"""
+JWT Authentication für Zen-AI-Pentest API
+"""
+
+from datetime import datetime, timedelta
+from typing import Optional, Dict
+from jose import JWTError, jwt
+from passlib.context import CryptContext
+from fastapi import Depends, HTTPException, status
+from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
+
+# Configuration
+SECRET_KEY = "your-secret-key-here-change-in-production"  # In production: env var
+ALGORITHM = "HS256"
+ACCESS_TOKEN_EXPIRE_MINUTES = 30
+
+# Password hashing
+pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")
+security = HTTPBearer()
+
+def verify_password(plain_password: str, hashed_password: str) -> bool:
+    """Verifiziert Passwort"""
+    return pwd_context.verify(plain_password, hashed_password)
+
+def get_password_hash(password: str) -> str:
+    """Hashed Passwort"""
+    return pwd_context.hash(password)
+
+def create_access_token(data: Dict, expires_delta: Optional[timedelta] = None) -> str:
+    """Erstellt JWT Token"""
+    to_encode = data.copy()
+    expire = datetime.utcnow() + (expires_delta or timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES))
+    to_encode.update({"exp": expire})
+    encoded_jwt = jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM)
+    return encoded_jwt
+
+def decode_token(token: str) -> Optional[Dict]:
+    """Decodiert JWT Token"""
+    try:
+        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
+        return payload
+    except JWTError:
+        return None
+
+async def verify_token(credentials: HTTPAuthorizationCredentials = Depends(security)) -> Dict:
+    """FastAPI Dependency für Token-Verifizierung"""
+    token = credentials.credentials
+    payload = decode_token(token)
+    
+    if payload is None:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Invalid authentication credentials",
+            headers={"WWW-Authenticate": "Bearer"},
+        )
+    
+    # Check expiration
+    exp = payload.get("exp")
+    if exp and datetime.utcnow().timestamp() > exp:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Token has expired",
+            headers={"WWW-Authenticate": "Bearer"},
+        )
+    
+    return payload
+
+def check_permissions(user: Dict, required_role: str) -> bool:
+    """Prüft ob User die benötigte Rolle hat"""
+    user_role = user.get("role", "viewer")
+    
+    roles_hierarchy = {
+        "viewer": 1,
+        "operator": 2,
+        "admin": 3
+    }
+    
+    user_level = roles_hierarchy.get(user_role, 0)
+    required_level = roles_hierarchy.get(required_role, 0)
+    
+    return user_level >= required_level
+
+async def require_admin(credentials: HTTPAuthorizationCredentials = Depends(security)) -> Dict:
+    """Erfordert Admin-Rolle"""
+    user = await verify_token(credentials)
+    
+    if not check_permissions(user, "admin"):
+        raise HTTPException(
+            status_code=status.HTTP_403_FORBIDDEN,
+            detail="Admin privileges required"
+        )
+    
+    return user
+
+async def require_operator(credentials: HTTPAuthorizationCredentials = Depends(security)) -> Dict:
+    """Erfordert Operator oder höhere Rolle"""
+    user = await verify_token(credentials)
+    
+    if not check_permissions(user, "operator"):
+        raise HTTPException(
+            status_code=status.HTTP_403_FORBIDDEN,
+            detail="Operator privileges required"
+        )
+    
+    return user
+
+# API Key Authentication (für CI/CD Integrationen)
+API_KEYS = {}  # In production: in DB speichern
+
+def verify_api_key(api_key: str) -> Optional[Dict]:
+    """Verifiziert API Key"""
+    return API_KEYS.get(api_key)
+
+def create_api_key(user_id: int, name: str) -> str:
+    """Erstellt neuen API Key"""
+    import secrets
+    key = secrets.token_urlsafe(32)
+    API_KEYS[key] = {
+        "user_id": user_id,
+        "name": name,
+        "created_at": datetime.utcnow().isoformat()
+    }
+    return key