aiptx 2.0.7__py3-none-any.whl

aipt_v2/tools/scanners/zap_tool.py

@@ -0,0 +1,612 @@
+"""
+OWASP ZAP (Zed Attack Proxy) Integration for AIPTX
+
+Provides web application security testing through ZAP's REST API.
+Supports spidering, active scanning, and passive scanning.
+
+Environment Variables:
+    ZAP_URL: ZAP API URL (default: http://localhost:8080)
+    ZAP_API_KEY: API key for authentication
+
+Usage:
+    from aipt_v2.tools.scanners.zap_tool import get_zap, zap_scan
+
+    zap = get_zap()
+    if zap.connect():
+        scan_id = zap.active_scan("https://example.com")
+        result = zap.wait_for_scan(scan_id)
+        alerts = zap.get_alerts()
+"""
+
+import json
+import logging
+import os
+import time
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from enum import Enum
+from typing import Any, Callable, Dict, List, Optional
+
+import requests
+
+logger = logging.getLogger(__name__)
+
+
+# ==================== Enums ====================
+
+class ScanStatus(Enum):
+    """ZAP scan status."""
+    RUNNING = "running"
+    COMPLETED = "completed"
+    STOPPED = "stopped"
+    FAILED = "failed"
+
+
+class RiskLevel(Enum):
+    """ZAP alert risk levels."""
+    HIGH = 3
+    MEDIUM = 2
+    LOW = 1
+    INFORMATIONAL = 0
+
+
+class ConfidenceLevel(Enum):
+    """ZAP alert confidence levels."""
+    HIGH = 3
+    MEDIUM = 2
+    LOW = 1
+    FALSE_POSITIVE = 0
+
+
+# ==================== Data Classes ====================
+
+@dataclass
+class ZAPConfig:
+    """Configuration for ZAP connection."""
+    base_url: str = field(default_factory=lambda: os.getenv("ZAP_URL", "http://localhost:8080"))
+    api_key: str = field(default_factory=lambda: os.getenv("ZAP_API_KEY", ""))
+    timeout: int = 30
+
+
+@dataclass
+class ScanResult:
+    """Result of a ZAP scan."""
+    scan_id: str
+    status: str
+    progress: int = 0
+    alerts_count: int = 0
+    urls_found: int = 0
+    start_time: str = ""
+    end_time: str = ""
+    error: str = ""
+
+
+@dataclass
+class Alert:
+    """ZAP security alert."""
+    alert_id: str
+    name: str
+    risk: int
+    risk_name: str
+    confidence: int
+    confidence_name: str
+    url: str
+    description: str = ""
+    solution: str = ""
+    reference: str = ""
+    cwe_id: int = 0
+    wasc_id: int = 0
+    evidence: str = ""
+    param: str = ""
+    attack: str = ""
+    other_info: str = ""
+
+
+# ==================== Main Tool Class ====================
+
+class ZAPTool:
+    """
+    OWASP ZAP integration for AIPTX.
+
+    Provides methods for:
+    - URL spidering
+    - Active vulnerability scanning
+    - Passive scanning
+    - Alert retrieval and management
+    """
+
+    def __init__(self, config: Optional[ZAPConfig] = None):
+        """Initialize ZAP tool with configuration."""
+        self.config = config or ZAPConfig()
+        self._session = requests.Session()
+        self._connected = False
+
+    def _request(self, component: str, action: str, params: dict = None) -> dict:
+        """Make request to ZAP API."""
+        url = f"{self.config.base_url}/JSON/{component}/view/{action}/"
+        params = params or {}
+        if self.config.api_key:
+            params["apikey"] = self.config.api_key
+
+        try:
+            response = self._session.get(url, params=params, timeout=self.config.timeout)
+            response.raise_for_status()
+            return response.json()
+        except requests.exceptions.RequestException as e:
+            logger.error(f"ZAP API error: {e}")
+            raise
+
+    def _action(self, component: str, action: str, params: dict = None) -> dict:
+        """Execute ZAP API action."""
+        url = f"{self.config.base_url}/JSON/{component}/action/{action}/"
+        params = params or {}
+        if self.config.api_key:
+            params["apikey"] = self.config.api_key
+
+        try:
+            response = self._session.get(url, params=params, timeout=self.config.timeout)
+            response.raise_for_status()
+            return response.json()
+        except requests.exceptions.RequestException as e:
+            logger.error(f"ZAP API action error: {e}")
+            raise
+
+    # ==================== Connection ====================
+
+    def connect(self) -> bool:
+        """Test connection to ZAP."""
+        try:
+            result = self._request("core", "version")
+            self._connected = "version" in result
+            if self._connected:
+                logger.info(f"Connected to ZAP at {self.config.base_url}")
+            return self._connected
+        except Exception as e:
+            logger.error(f"Failed to connect to ZAP: {e}")
+            self._connected = False
+            return False
+
+    def is_connected(self) -> bool:
+        """Check if connected to ZAP."""
+        return self._connected
+
+    def get_info(self) -> Dict:
+        """Get ZAP version and info."""
+        try:
+            version = self._request("core", "version")
+            return {"version": version.get("version", "unknown")}
+        except Exception:
+            return {}
+
+    # ==================== Context Management ====================
+
+    def new_session(self, name: str = None) -> bool:
+        """Create a new ZAP session."""
+        try:
+            params = {}
+            if name:
+                params["name"] = name
+            self._action("core", "newSession", params)
+            return True
+        except Exception:
+            return False
+
+    def set_mode(self, mode: str = "standard") -> bool:
+        """Set ZAP mode (safe, protect, standard, attack)."""
+        try:
+            self._action("core", "setMode", {"mode": mode})
+            return True
+        except Exception:
+            return False
+
+    # ==================== Spider ====================
+
+    def spider(self, url: str, max_children: int = 0, recurse: bool = True) -> str:
+        """
+        Start spidering a URL.
+
+        Args:
+            url: Target URL to spider
+            max_children: Max children to crawl (0 = unlimited)
+            recurse: Whether to recurse into found URLs
+
+        Returns:
+            Spider scan ID
+        """
+        try:
+            params = {
+                "url": url,
+                "maxChildren": str(max_children),
+                "recurse": str(recurse).lower()
+            }
+            result = self._action("spider", "scan", params)
+            scan_id = result.get("scan", "")
+            logger.info(f"Started spider scan: {scan_id} for {url}")
+            return str(scan_id)
+        except Exception as e:
+            logger.error(f"Failed to start spider: {e}")
+            return ""
+
+    def get_spider_status(self, scan_id: str) -> int:
+        """Get spider progress (0-100)."""
+        try:
+            result = self._request("spider", "status", {"scanId": scan_id})
+            return int(result.get("status", 0))
+        except Exception:
+            return 0
+
+    def wait_for_spider(self, scan_id: str, timeout: int = 300) -> bool:
+        """Wait for spider to complete."""
+        start_time = time.time()
+        while time.time() - start_time < timeout:
+            progress = self.get_spider_status(scan_id)
+            if progress >= 100:
+                return True
+            time.sleep(5)
+        return False
+
+    def get_spider_results(self, scan_id: str) -> List[str]:
+        """Get URLs found by spider."""
+        try:
+            result = self._request("spider", "results", {"scanId": scan_id})
+            return result.get("results", [])
+        except Exception:
+            return []
+
+    # ==================== Active Scan ====================
+
+    def active_scan(
+        self,
+        url: str,
+        recurse: bool = True,
+        in_scope_only: bool = False,
+        scan_policy: str = None
+    ) -> str:
+        """
+        Start active vulnerability scan.
+
+        Args:
+            url: Target URL to scan
+            recurse: Scan recursively
+            in_scope_only: Only scan URLs in scope
+            scan_policy: Scan policy name (optional)
+
+        Returns:
+            Active scan ID
+        """
+        try:
+            params = {
+                "url": url,
+                "recurse": str(recurse).lower(),
+                "inScopeOnly": str(in_scope_only).lower()
+            }
+            if scan_policy:
+                params["scanPolicyName"] = scan_policy
+
+            result = self._action("ascan", "scan", params)
+            scan_id = result.get("scan", "")
+            logger.info(f"Started active scan: {scan_id} for {url}")
+            return str(scan_id)
+        except Exception as e:
+            logger.error(f"Failed to start active scan: {e}")
+            return ""
+
+    def get_scan_status(self, scan_id: str) -> ScanResult:
+        """Get active scan status."""
+        try:
+            status_result = self._request("ascan", "status", {"scanId": scan_id})
+            progress = int(status_result.get("status", 0))
+
+            # Get alerts count
+            alerts = self.get_alerts()
+
+            status = "running" if progress < 100 else "completed"
+
+            return ScanResult(
+                scan_id=scan_id,
+                status=status,
+                progress=progress,
+                alerts_count=len(alerts)
+            )
+        except Exception as e:
+            return ScanResult(scan_id=scan_id, status="error", error=str(e))
+
+    def wait_for_scan(
+        self,
+        scan_id: str,
+        timeout: int = 3600,
+        poll_interval: int = 10,
+        callback: Callable[[ScanResult], None] = None
+    ) -> ScanResult:
+        """
+        Wait for active scan to complete.
+
+        Args:
+            scan_id: Scan ID to monitor
+            timeout: Maximum wait time in seconds
+            poll_interval: Polling interval in seconds
+            callback: Optional progress callback
+
+        Returns:
+            Final scan result
+        """
+        start_time = time.time()
+        while time.time() - start_time < timeout:
+            result = self.get_scan_status(scan_id)
+
+            if callback:
+                callback(result)
+
+            if result.progress >= 100 or result.status in ["completed", "stopped", "error"]:
+                result.status = "completed" if result.progress >= 100 else result.status
+                return result
+
+            time.sleep(poll_interval)
+
+        return ScanResult(
+            scan_id=scan_id,
+            status="timeout",
+            error=f"Scan timed out after {timeout}s"
+        )
+
+    def stop_scan(self, scan_id: str) -> bool:
+        """Stop an active scan."""
+        try:
+            self._action("ascan", "stop", {"scanId": scan_id})
+            return True
+        except Exception:
+            return False
+
+    def pause_scan(self, scan_id: str) -> bool:
+        """Pause an active scan."""
+        try:
+            self._action("ascan", "pause", {"scanId": scan_id})
+            return True
+        except Exception:
+            return False
+
+    def resume_scan(self, scan_id: str) -> bool:
+        """Resume a paused scan."""
+        try:
+            self._action("ascan", "resume", {"scanId": scan_id})
+            return True
+        except Exception:
+            return False
+
+    # ==================== Full Scan (Spider + Active) ====================
+
+    def full_scan(self, url: str, spider_timeout: int = 300) -> str:
+        """
+        Run a full scan (spider + active scan).
+
+        Args:
+            url: Target URL
+            spider_timeout: Max time for spider phase
+
+        Returns:
+            Active scan ID
+        """
+        # First spider the target
+        spider_id = self.spider(url)
+        if spider_id:
+            self.wait_for_spider(spider_id, spider_timeout)
+
+        # Then run active scan
+        return self.active_scan(url)
+
+    # ==================== Alert Management ====================
+
+    def get_alerts(
+        self,
+        base_url: str = None,
+        risk: str = None,
+        start: int = 0,
+        count: int = 1000
+    ) -> List[Alert]:
+        """
+        Get security alerts.
+
+        Args:
+            base_url: Filter by base URL
+            risk: Filter by risk level (high, medium, low, informational)
+            start: Start index
+            count: Number of alerts to retrieve
+
+        Returns:
+            List of Alert objects
+        """
+        try:
+            params = {"start": str(start), "count": str(count)}
+            if base_url:
+                params["baseurl"] = base_url
+            if risk:
+                params["riskId"] = str(self._risk_to_int(risk))
+
+            result = self._request("alert", "alerts", params)
+            alerts = []
+
+            for a in result.get("alerts", []):
+                alerts.append(Alert(
+                    alert_id=str(a.get("id", "")),
+                    name=a.get("name", "Unknown"),
+                    risk=int(a.get("riskcode", 0)),
+                    risk_name=a.get("risk", "Unknown"),
+                    confidence=int(a.get("confidence", 0)),
+                    confidence_name=a.get("confidence", "Unknown"),
+                    url=a.get("url", ""),
+                    description=a.get("description", ""),
+                    solution=a.get("solution", ""),
+                    reference=a.get("reference", ""),
+                    cwe_id=int(a.get("cweid", 0) or 0),
+                    wasc_id=int(a.get("wascid", 0) or 0),
+                    evidence=a.get("evidence", ""),
+                    param=a.get("param", ""),
+                    attack=a.get("attack", ""),
+                    other_info=a.get("other", "")
+                ))
+
+            return alerts
+        except Exception as e:
+            logger.error(f"Failed to get alerts: {e}")
+            return []
+
+    def _risk_to_int(self, risk: str) -> int:
+        """Convert risk string to int."""
+        mapping = {"high": 3, "medium": 2, "low": 1, "informational": 0, "info": 0}
+        return mapping.get(risk.lower(), 0)
+
+    def get_alerts_summary(self) -> Dict:
+        """Get summary of alerts by risk level."""
+        try:
+            result = self._request("alert", "alertsSummary")
+            return result.get("alertsSummary", {})
+        except Exception:
+            return {}
+
+    # ==================== AIPTX Integration ====================
+
+    def to_finding(self, alert: Alert) -> Dict:
+        """Convert ZAP alert to AIPTX finding format."""
+        severity_map = {3: "high", 2: "medium", 1: "low", 0: "info"}
+
+        return {
+            "type": "vulnerability",
+            "value": alert.name,
+            "description": alert.description or alert.name,
+            "severity": severity_map.get(alert.risk, "info"),
+            "phase": "scan",
+            "tool": "zap",
+            "target": alert.url,
+            "metadata": {
+                "alert_id": alert.alert_id,
+                "risk": alert.risk_name,
+                "confidence": alert.confidence_name,
+                "cwe_id": alert.cwe_id,
+                "wasc_id": alert.wasc_id,
+                "solution": alert.solution,
+                "evidence": alert.evidence[:500] if alert.evidence else "",
+                "param": alert.param,
+                "attack": alert.attack
+            },
+            "timestamp": datetime.now(timezone.utc).isoformat()
+        }
+
+    def export_findings(self, output_path: str = None, base_url: str = None) -> List[Dict]:
+        """
+        Export alerts in AIPTX format.
+
+        Args:
+            output_path: Optional file path to save JSON
+            base_url: Optional URL filter
+
+        Returns:
+            List of AIPTX findings
+        """
+        alerts = self.get_alerts(base_url=base_url)
+        findings = [self.to_finding(a) for a in alerts]
+
+        if output_path:
+            with open(output_path, "w") as f:
+                json.dump(findings, f, indent=2)
+
+        return findings
+
+
+# ==================== Global Instance & Helper Functions ====================
+
+_zap: Optional[ZAPTool] = None
+
+
+def get_zap(config: Optional[ZAPConfig] = None) -> ZAPTool:
+    """Get or create global ZAP instance."""
+    global _zap
+    if _zap is None or config is not None:
+        _zap = ZAPTool(config)
+    return _zap
+
+
+def zap_scan(url: str, full: bool = True) -> ScanResult:
+    """Quick scan a target."""
+    zap = get_zap()
+    if not zap.connect():
+        return ScanResult(scan_id="", status="error", error="Connection failed")
+
+    if full:
+        scan_id = zap.full_scan(url)
+    else:
+        scan_id = zap.active_scan(url)
+
+    return zap.get_scan_status(scan_id)
+
+
+def zap_status(scan_id: str) -> ScanResult:
+    """Get scan status."""
+    return get_zap().get_scan_status(scan_id)
+
+
+def zap_alerts(base_url: str = None, risk: str = None) -> List[Dict]:
+    """Get alerts in AIPTX format."""
+    zap = get_zap()
+    alerts = zap.get_alerts(base_url=base_url, risk=risk)
+    return [zap.to_finding(a) for a in alerts]
+
+
+def zap_summary() -> Dict:
+    """Get ZAP scanner summary."""
+    zap = get_zap()
+    if not zap.connect():
+        return {"connected": False, "error": "Connection failed"}
+
+    info = zap.get_info()
+    alerts_summary = zap.get_alerts_summary()
+
+    return {
+        "connected": True,
+        "url": zap.config.base_url,
+        "version": info.get("version", "unknown"),
+        "alerts": alerts_summary
+    }
+
+
+# ==================== CLI Testing ====================
+
+if __name__ == "__main__":
+    import sys
+
+    print("OWASP ZAP Integration Test")
+    print("=" * 50)
+
+    zap = get_zap()
+
+    if zap.connect():
+        print(f"✓ Connected to ZAP at {zap.config.base_url}")
+
+        info = zap.get_info()
+        print(f"  Version: {info.get('version', 'unknown')}")
+
+        if len(sys.argv) > 1:
+            target = sys.argv[1]
+            print(f"\nScanning {target}...")
+
+            # Spider first
+            spider_id = zap.spider(target)
+            print(f"  Spider ID: {spider_id}")
+            zap.wait_for_spider(spider_id)
+            urls = zap.get_spider_results(spider_id)
+            print(f"  URLs found: {len(urls)}")
+
+            # Active scan
+            scan_id = zap.active_scan(target)
+            print(f"  Scan ID: {scan_id}")
+
+            result = zap.wait_for_scan(scan_id, timeout=1800)
+            print(f"  Status: {result.status}")
+
+            alerts = zap.get_alerts()
+            print(f"  Alerts: {len(alerts)}")
+
+            for alert in alerts[:5]:
+                print(f"    [{alert.risk_name}] {alert.name}")
+    else:
+        print("✗ Failed to connect to ZAP")
+        print("  Set ZAP_URL and ZAP_API_KEY environment variables")

aipt_v2/tools/terminal/__init__.py

@@ -0,0 +1,5 @@
+"""
+AIPT Terminal Tools - Shell session management
+"""
+
+__all__ = []

aipt_v2/tools/terminal/terminal_actions.py

@@ -0,0 +1,37 @@
+from __future__ import annotations
+
+from typing import Any
+
+from aipt_v2.tools.registry import register_tool
+
+from .terminal_manager import get_terminal_manager
+
+
+@register_tool
+def terminal_execute(
+    command: str,
+    is_input: bool = False,
+    timeout: float | None = None,
+    terminal_id: str | None = None,
+    no_enter: bool = False,
+) -> dict[str, Any]:
+    manager = get_terminal_manager()
+
+    try:
+        return manager.execute_command(
+            command=command,
+            is_input=is_input,
+            timeout=timeout,
+            terminal_id=terminal_id,
+            no_enter=no_enter,
+        )
+    except (ValueError, RuntimeError) as e:
+        return {
+            "error": str(e),
+            "command": command,
+            "terminal_id": terminal_id or "default",
+            "content": "",
+            "status": "error",
+            "exit_code": None,
+            "working_dir": None,
+        }