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

api/schemas.py

@@ -0,0 +1,357 @@
+"""
+Pydantic Schemas für API Validierung
+"""
+
+from pydantic import BaseModel, Field
+from typing import Optional, List, Dict, Any
+from datetime import datetime
+from enum import Enum
+
+# ============================================================================
+# ENUMS
+# ============================================================================
+
+class ScanStatus(str, Enum):
+    PENDING = "pending"
+    RUNNING = "running"
+    COMPLETED = "completed"
+    FAILED = "failed"
+    CANCELLED = "cancelled"
+
+class Severity(str, Enum):
+    CRITICAL = "critical"
+    HIGH = "high"
+    MEDIUM = "medium"
+    LOW = "low"
+    INFO = "info"
+
+class ReportFormat(str, Enum):
+    PDF = "pdf"
+    HTML = "html"
+    JSON = "json"
+    XML = "xml"
+
+# ============================================================================
+# BASE SCHEMAS
+# ============================================================================
+
+class ScanBase(BaseModel):
+    name: str = Field(..., min_length=1, max_length=255)
+    target: str = Field(..., min_length=1, max_length=500)
+    scan_type: str = Field(..., min_length=1, max_length=100)
+    config: Optional[Dict[str, Any]] = {}
+
+class ScanCreate(ScanBase):
+    objective: Optional[str] = "comprehensive security scan"
+
+class ScanUpdate(BaseModel):
+    status: Optional[str] = None
+    config: Optional[Dict[str, Any]] = None
+
+class ScanResponse(ScanBase):
+    id: int
+    status: str
+    user_id: int
+    created_at: datetime
+    started_at: Optional[datetime] = None
+    completed_at: Optional[datetime] = None
+    result_summary: Optional[str] = None
+    findings_count: Optional[int] = 0
+    
+    class Config:
+        from_attributes = True
+
+# ============================================================================
+# FINDING SCHEMAS
+# ============================================================================
+
+class FindingBase(BaseModel):
+    title: str = Field(..., min_length=1, max_length=500)
+    description: Optional[str] = None
+    severity: Severity = Severity.MEDIUM
+    cvss_score: Optional[float] = Field(None, ge=0, le=10)
+    cve_id: Optional[str] = None
+    evidence: Optional[str] = None
+    remediation: Optional[str] = None
+    tool: Optional[str] = None
+    target: Optional[str] = None
+    port: Optional[int] = None
+    service: Optional[str] = None
+
+class FindingCreate(FindingBase):
+    pass
+
+class FindingResponse(FindingBase):
+    id: int
+    scan_id: int
+    created_at: datetime
+    verified: int = 0
+    
+    class Config:
+        from_attributes = True
+
+class FindingUpdate(BaseModel):
+    severity: Optional[Severity] = None
+    verified: Optional[int] = None
+    remediation: Optional[str] = None
+
+# ============================================================================
+# REPORT SCHEMAS
+# ============================================================================
+
+class ReportBase(BaseModel):
+    scan_id: int
+    format: ReportFormat = ReportFormat.PDF
+    template: Optional[str] = "default"
+
+class ReportCreate(ReportBase):
+    pass
+
+class ReportResponse(ReportBase):
+    id: int
+    user_id: int
+    status: str
+    file_path: Optional[str] = None
+    file_size: Optional[int] = None
+    created_at: datetime
+    generated_at: Optional[datetime] = None
+    
+    class Config:
+        from_attributes = True
+
+# ============================================================================
+# TOOL EXECUTION SCHEMAS
+# ============================================================================
+
+class ToolExecuteRequest(BaseModel):
+    tool_name: str = Field(..., description="Name of the tool to execute")
+    target: str = Field(..., description="Target to scan")
+    parameters: Optional[Dict[str, Any]] = {}
+    timeout: Optional[int] = 300
+
+class ToolExecuteResponse(BaseModel):
+    scan_id: int
+    status: str
+    message: str
+    estimated_duration: Optional[int] = None
+
+class ToolInfo(BaseModel):
+    name: str
+    description: str
+    category: str
+    parameters: Optional[Dict[str, Any]] = {}
+
+# ============================================================================
+# AUTH SCHEMAS
+# ============================================================================
+
+class UserLogin(BaseModel):
+    username: str
+    password: str
+
+class UserCreate(BaseModel):
+    username: str = Field(..., min_length=3, max_length=100)
+    email: str = Field(..., pattern=r'^[\w\.-]+@[\w\.-]+\.\w+$')
+    password: str = Field(..., min_length=8)
+    role: Optional[str] = "operator"
+
+class UserResponse(BaseModel):
+    id: int
+    username: str
+    email: str
+    role: str
+    is_active: int
+    created_at: datetime
+    
+    class Config:
+        from_attributes = True
+
+class TokenResponse(BaseModel):
+    access_token: str
+    token_type: str = "bearer"
+    expires_in: int = 3600
+
+# ============================================================================
+# DASHBOARD SCHEMAS
+# ============================================================================
+
+class DashboardStats(BaseModel):
+    total_scans: int
+    active_scans: int
+    completed_scans: int
+    failed_scans: int
+    total_findings: int
+    critical_findings: int
+    high_findings: int
+    reports_generated: int
+
+class RecentActivity(BaseModel):
+    id: int
+    type: str  # scan, finding, report
+    description: str
+    timestamp: datetime
+    user: str
+
+class DashboardResponse(BaseModel):
+    stats: DashboardStats
+    recent_activities: List[RecentActivity]
+    scans_by_status: Dict[str, int]
+    findings_by_severity: Dict[str, int]
+
+# ============================================================================
+# WEBSOCKET SCHEMAS
+# ============================================================================
+
+class WSMessage(BaseModel):
+    type: str  # status, log, result, error
+    scan_id: Optional[int] = None
+    message: str
+    data: Optional[Dict[str, Any]] = None
+    timestamp: datetime = Field(default_factory=datetime.utcnow)
+
+class WSCommand(BaseModel):
+    action: str  # start, stop, status, ping
+    scan_id: Optional[int] = None
+    parameters: Optional[Dict[str, Any]] = None
+
+# ============================================================================
+# NOTIFICATION SCHEMAS
+# ============================================================================
+
+class NotificationBase(BaseModel):
+    type: str
+    title: str
+    message: str
+
+class NotificationCreate(NotificationBase):
+    user_id: int
+
+class NotificationResponse(NotificationBase):
+    id: int
+    user_id: int
+    read: int
+    created_at: datetime
+    
+    class Config:
+        from_attributes = True
+
+# ============================================================================
+# ASSET SCHEMAS
+# ============================================================================
+
+class AssetBase(BaseModel):
+    name: str
+    asset_type: str
+    ip_address: Optional[str] = None
+    hostname: Optional[str] = None
+    os: Optional[str] = None
+    criticality: str = "medium"
+
+class AssetCreate(AssetBase):
+    pass
+
+class AssetResponse(AssetBase):
+    id: int
+    services: Optional[Dict[str, Any]] = None
+    created_at: datetime
+    last_scanned: Optional[datetime] = None
+    
+    class Config:
+        from_attributes = True
+
+# ============================================================================
+# VULNERABILITY DB SCHEMAS
+# ============================================================================
+
+class VulnerabilityDBResponse(BaseModel):
+    id: int
+    cve_id: str
+    title: str
+    description: Optional[str] = None
+    severity: str
+    cvss_score: Optional[float] = None
+    epss_score: Optional[float] = None
+    affected_products: Optional[List[str]] = None
+    references: Optional[List[str]] = None
+    exploits: Optional[List[str]] = None
+
+# ============================================================================
+# PAGINATION SCHEMAS
+# ============================================================================
+
+class PaginatedResponse(BaseModel):
+    items: List[Any]
+    total: int
+    page: int
+    page_size: int
+    pages: int
+
+class PaginationParams(BaseModel):
+    page: int = Field(1, ge=1)
+    page_size: int = Field(20, ge=1, le=100)
+    
+    @property
+    def skip(self) -> int:
+        return (self.page - 1) * self.page_size
+
+# ============================================================================
+# SCHEDULED SCAN SCHEMAS
+# ============================================================================
+
+class ScheduleFrequency(str, Enum):
+    ONCE = "once"
+    DAILY = "daily"
+    WEEKLY = "weekly"
+    MONTHLY = "monthly"
+
+class ScheduledScanCreate(BaseModel):
+    name: str = Field(..., min_length=1, max_length=200)
+    target: str = Field(..., min_length=1, max_length=500)
+    scan_type: str = "comprehensive"
+    frequency: ScheduleFrequency = ScheduleFrequency.WEEKLY
+    schedule_time: str = Field(..., pattern=r"^([0-1]?[0-9]|2[0-3]):[0-5][0-9]$")  # HH:MM format
+    schedule_day: Optional[int] = Field(None, ge=0, le=6)  # 0=Monday, 6=Sunday for weekly
+    enabled: bool = True
+    notification_email: Optional[str] = None
+    notification_slack: Optional[str] = None
+
+class ScheduledScanUpdate(BaseModel):
+    name: Optional[str] = None
+    target: Optional[str] = None
+    scan_type: Optional[str] = None
+    frequency: Optional[ScheduleFrequency] = None
+    schedule_time: Optional[str] = None
+    schedule_day: Optional[int] = None
+    enabled: Optional[bool] = None
+    notification_email: Optional[str] = None
+    notification_slack: Optional[str] = None
+    last_run_status: Optional[str] = None
+
+class ScheduledScanResponse(BaseModel):
+    id: int
+    name: str
+    target: str
+    scan_type: str
+    frequency: str
+    schedule_time: str
+    schedule_day: Optional[int] = None
+    enabled: bool
+    notification_email: Optional[str] = None
+    notification_slack: Optional[str] = None
+    last_run_at: Optional[datetime] = None
+    last_run_status: Optional[str] = None
+    next_run_at: Optional[datetime] = None
+    created_at: datetime
+    created_by: str
+    
+    class Config:
+        from_attributes = True
+
+class ScheduleExecutionResponse(BaseModel):
+    id: int
+    scheduled_scan_id: int
+    scan_id: Optional[int] = None
+    status: str
+    started_at: Optional[datetime] = None
+    completed_at: Optional[datetime] = None
+    error_message: Optional[str] = None

api/websocket.py

@@ -0,0 +1,97 @@
+"""
+WebSocket Connection Manager für Real-Time Updates
+"""
+
+import json
+import logging
+from typing import Dict, List, Set
+from fastapi import WebSocket
+
+logger = logging.getLogger(__name__)
+
+class ConnectionManager:
+    """Manages WebSocket connections for real-time updates"""
+    
+    def __init__(self):
+        # scan_id -> Set of WebSocket connections
+        self.scan_connections: Dict[int, Set[WebSocket]] = {}
+        # Global connections (for notifications)
+        self.global_connections: Set[WebSocket] = set()
+    
+    async def connect(self, websocket: WebSocket, scan_id: int = None):
+        """Accept and register new connection"""
+        await websocket.accept()
+        
+        if scan_id:
+            if scan_id not in self.scan_connections:
+                self.scan_connections[scan_id] = set()
+            self.scan_connections[scan_id].add(websocket)
+            logger.info(f"WebSocket connected for scan {scan_id}")
+        else:
+            self.global_connections.add(websocket)
+            logger.info("Global WebSocket connected")
+    
+    def disconnect(self, websocket: WebSocket, scan_id: int = None):
+        """Remove connection"""
+        if scan_id and scan_id in self.scan_connections:
+            self.scan_connections[scan_id].discard(websocket)
+            if not self.scan_connections[scan_id]:
+                del self.scan_connections[scan_id]
+        else:
+            self.global_connections.discard(websocket)
+        
+        logger.info(f"WebSocket disconnected (scan: {scan_id})")
+    
+    async def broadcast_to_scan(self, scan_id: int, message: dict):
+        """Send message to all connections for a scan"""
+        if scan_id not in self.scan_connections:
+            return
+        
+        disconnected = set()
+        for connection in self.scan_connections[scan_id]:
+            try:
+                await connection.send_json(message)
+            except Exception as e:
+                logger.error(f"WebSocket send error: {e}")
+                disconnected.add(connection)
+        
+        # Clean up disconnected
+        for conn in disconnected:
+            self.disconnect(conn, scan_id)
+    
+    async def broadcast_global(self, message: dict):
+        """Send message to all global connections"""
+        disconnected = set()
+        for connection in self.global_connections:
+            try:
+                await connection.send_json(message)
+            except Exception as e:
+                logger.error(f"WebSocket send error: {e}")
+                disconnected.add(connection)
+        
+        # Clean up disconnected
+        for conn in disconnected:
+            self.disconnect(conn)
+    
+    async def send_personal_message(self, websocket: WebSocket, message: dict):
+        """Send message to specific connection"""
+        try:
+            await websocket.send_json(message)
+        except Exception as e:
+            logger.error(f"WebSocket send error: {e}")
+    
+    def get_scan_connections_count(self, scan_id: int) -> int:
+        """Get number of active connections for a scan"""
+        return len(self.scan_connections.get(scan_id, set()))
+    
+    def get_global_connections_count(self) -> int:
+        """Get number of global connections"""
+        return len(self.global_connections)
+    
+    def get_stats(self) -> dict:
+        """Get connection statistics"""
+        return {
+            "global_connections": len(self.global_connections),
+            "active_scans": len(self.scan_connections),
+            "total_scan_connections": sum(len(conns) for conns in self.scan_connections.values())
+        }

autonomous/__init__.py

@@ -0,0 +1,122 @@
+"""
+Autonomous Agent System for Zen AI Pentest (2026 Roadmap - Q1)
+
+This module implements:
+- ReAct/Plan-and-Execute reasoning loop
+- Real tool execution framework
+- Memory system with vector storage
+- Self-correction capabilities
+- Exploit validation with sandboxed execution
+
+Usage:
+    from autonomous import AutonomousAgent
+    
+    agent = AutonomousAgent(goal="Find RCE in target API")
+    result = await agent.run()
+
+    # Exploit Validation
+    from autonomous import ExploitValidator, ExploitType
+    
+    validator = ExploitValidator()
+    result = await validator.validate(
+        exploit_code="...",
+        target="https://example.com",
+        exploit_type=ExploitType.WEB_SQLI
+    )
+
+    # Autonomous Agent Loop
+    from autonomous import AutonomousAgentLoop, AgentState, AgentMemory
+    
+    agent_loop = AutonomousAgentLoop(max_iterations=50)
+    result = await agent_loop.run(
+        goal="Find vulnerabilities",
+        target="example.com",
+        scope={"depth": "comprehensive"}
+    )
+"""
+
+# Core autonomous components
+from .agent import AutonomousAgent
+from .react import ReActLoop, Thought, Action, Observation
+from .tool_executor import ToolExecutor, ToolRegistry
+from .memory import MemoryManager, WorkingMemory, LongTermMemory
+
+# Exploit Validator components
+from .exploit_validator import (
+    ExploitValidator,
+    ExploitValidatorPool,
+    ExploitResult,
+    ExploitType,
+    ExploitStatus,
+    SafetyLevel,
+    Evidence,
+    ScopeConfig,
+    SandboxConfig,
+    validate_sql_injection,
+    validate_xss,
+    validate_command_injection,
+)
+
+# Autonomous Agent Loop components
+from .agent_loop import (
+    AutonomousAgentLoop,
+    AgentState,
+    AgentMemory,
+    ToolResult,
+    PlanStep,
+    ToolType,
+    BaseTool,
+    NmapScanner,
+    NucleiScanner,
+    ExploitValidator as ExploitValidatorTool,
+    ReportGenerator,
+    SubdomainEnumerator,
+    ToolRegistry as AgentToolRegistry,
+    create_agent_loop
+)
+
+__all__ = [
+    # Core autonomous components
+    'AutonomousAgent',
+    'ReActLoop',
+    'Thought',
+    'Action',
+    'Observation',
+    'ToolExecutor',
+    'ToolRegistry',
+    'MemoryManager',
+    'WorkingMemory',
+    'LongTermMemory',
+    
+    # Exploit Validator components
+    'ExploitValidator',
+    'ExploitValidatorPool',
+    'ExploitResult',
+    'ExploitType',
+    'ExploitStatus',
+    'SafetyLevel',
+    'Evidence',
+    'ScopeConfig',
+    'SandboxConfig',
+    'validate_sql_injection',
+    'validate_xss',
+    'validate_command_injection',
+    
+    # Autonomous Agent Loop components
+    'AutonomousAgentLoop',
+    'AgentState',
+    'AgentMemory',
+    'ToolResult',
+    'PlanStep',
+    'ToolType',
+    'BaseTool',
+    'NmapScanner',
+    'NucleiScanner',
+    'ExploitValidatorTool',
+    'ReportGenerator',
+    'SubdomainEnumerator',
+    'AgentToolRegistry',
+    'create_agent_loop',
+]
+
+__version__ = '2.0.0'