aiptx 2.0.2__py3-none-any.whl

aipt_v2/tools/registry.py

@@ -0,0 +1,196 @@
+import inspect
+import logging
+import os
+from collections.abc import Callable
+from functools import wraps
+from inspect import signature
+from pathlib import Path
+from typing import Any
+
+
+tools: list[dict[str, Any]] = []
+_tools_by_name: dict[str, Callable[..., Any]] = {}
+logger = logging.getLogger(__name__)
+
+
+class ImplementedInClientSideOnlyError(Exception):
+    def __init__(
+        self,
+        message: str = "This tool is implemented in the client side only",
+    ) -> None:
+        self.message = message
+        super().__init__(self.message)
+
+
+def _process_dynamic_content(content: str) -> str:
+    if "{{DYNAMIC_MODULES_DESCRIPTION}}" in content:
+        try:
+            from prompts import generate_modules_description
+
+            modules_description = generate_modules_description()
+            content = content.replace("{{DYNAMIC_MODULES_DESCRIPTION}}", modules_description)
+        except ImportError:
+            logger.warning("Could not import prompts utilities for dynamic schema generation")
+            content = content.replace(
+                "{{DYNAMIC_MODULES_DESCRIPTION}}",
+                "List of prompt modules to load for this agent (max 5). Module discovery failed.",
+            )
+
+    return content
+
+
+def _load_xml_schema(path: Path) -> Any:
+    if not path.exists():
+        return None
+    try:
+        content = path.read_text()
+
+        content = _process_dynamic_content(content)
+
+        start_tag = '<tool name="'
+        end_tag = "</tool>"
+        tools_dict = {}
+
+        pos = 0
+        while True:
+            start_pos = content.find(start_tag, pos)
+            if start_pos == -1:
+                break
+
+            name_start = start_pos + len(start_tag)
+            name_end = content.find('"', name_start)
+            if name_end == -1:
+                break
+            tool_name = content[name_start:name_end]
+
+            end_pos = content.find(end_tag, name_end)
+            if end_pos == -1:
+                break
+            end_pos += len(end_tag)
+
+            tool_element = content[start_pos:end_pos]
+            tools_dict[tool_name] = tool_element
+
+            pos = end_pos
+
+            if pos >= len(content):
+                break
+    except (IndexError, ValueError, UnicodeError) as e:
+        logger.warning(f"Error loading schema file {path}: {e}")
+        return None
+    else:
+        return tools_dict
+
+
+def _get_module_name(func: Callable[..., Any]) -> str:
+    module = inspect.getmodule(func)
+    if not module:
+        return "unknown"
+
+    module_name = module.__name__
+    if ".tools." in module_name:
+        parts = module_name.split(".tools.")[-1].split(".")
+        if len(parts) >= 1:
+            return parts[0]
+    return "unknown"
+
+
+def register_tool(
+    func: Callable[..., Any] | None = None, *, sandbox_execution: bool = True
+) -> Callable[..., Any]:
+    def decorator(f: Callable[..., Any]) -> Callable[..., Any]:
+        func_dict = {
+            "name": f.__name__,
+            "function": f,
+            "module": _get_module_name(f),
+            "sandbox_execution": sandbox_execution,
+        }
+
+        sandbox_mode = os.getenv("AIPT_SANDBOX_MODE", "false").lower() == "true"
+        if not sandbox_mode:
+            try:
+                module_path = Path(inspect.getfile(f))
+                schema_file_name = f"{module_path.stem}_schema.xml"
+                schema_path = module_path.parent / schema_file_name
+
+                xml_tools = _load_xml_schema(schema_path)
+
+                if xml_tools is not None and f.__name__ in xml_tools:
+                    func_dict["xml_schema"] = xml_tools[f.__name__]
+                else:
+                    func_dict["xml_schema"] = (
+                        f'<tool name="{f.__name__}">'
+                        "<description>Schema not found for tool.</description>"
+                        "</tool>"
+                    )
+            except (TypeError, FileNotFoundError) as e:
+                logger.warning(f"Error loading schema for {f.__name__}: {e}")
+                func_dict["xml_schema"] = (
+                    f'<tool name="{f.__name__}">'
+                    "<description>Error loading schema.</description>"
+                    "</tool>"
+                )
+
+        tools.append(func_dict)
+        _tools_by_name[str(func_dict["name"])] = f
+
+        @wraps(f)
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            return f(*args, **kwargs)
+
+        return wrapper
+
+    if func is None:
+        return decorator
+    return decorator(func)
+
+
+def get_tool_by_name(name: str) -> Callable[..., Any] | None:
+    return _tools_by_name.get(name)
+
+
+def get_tool_names() -> list[str]:
+    return list(_tools_by_name.keys())
+
+
+def needs_agent_state(tool_name: str) -> bool:
+    tool_func = get_tool_by_name(tool_name)
+    if not tool_func:
+        return False
+    sig = signature(tool_func)
+    return "agent_state" in sig.parameters
+
+
+def should_execute_in_sandbox(tool_name: str) -> bool:
+    for tool in tools:
+        if tool.get("name") == tool_name:
+            return bool(tool.get("sandbox_execution", True))
+    return True
+
+
+def get_tools_prompt() -> str:
+    tools_by_module: dict[str, list[dict[str, Any]]] = {}
+    for tool in tools:
+        module = tool.get("module", "unknown")
+        if module not in tools_by_module:
+            tools_by_module[module] = []
+        tools_by_module[module].append(tool)
+
+    xml_sections = []
+    for module, module_tools in sorted(tools_by_module.items()):
+        tag_name = f"{module}_tools"
+        section_parts = [f"<{tag_name}>"]
+        for tool in module_tools:
+            tool_xml = tool.get("xml_schema", "")
+            if tool_xml:
+                indented_tool = "\n".join(f"  {line}" for line in tool_xml.split("\n"))
+                section_parts.append(indented_tool)
+        section_parts.append(f"</{tag_name}>")
+        xml_sections.append("\n".join(section_parts))
+
+    return "\n\n".join(xml_sections)
+
+
+def clear_registry() -> None:
+    tools.clear()
+    _tools_by_name.clear()

aipt_v2/tools/scanners/__init__.py

@@ -0,0 +1,343 @@
+"""
+AIPT Scanner Tools - External Scanner Integrations
+
+Provides plug-and-play integrations with enterprise security scanners:
+- Acunetix Web Vulnerability Scanner
+- Burp Suite Pro/Enterprise
+- Nessus Professional/Expert (Network Vulnerability Scanner)
+- OWASP ZAP (Dynamic Application Security Testing)
+
+Server Configuration:
+    Public IP:    13.127.28.41
+    Acunetix:     https://13.127.28.41:3443
+    Burp Suite:   http://13.127.28.41:1337
+    Nessus:       https://your-nessus-server:8834
+    ZAP:          http://localhost:8080
+"""
+
+# Configuration
+from aipt_v2.tools.scanners.config import (
+    SCANNER_SERVER_IP,
+    ACUNETIX_PORT,
+    BURP_PORT,
+    ACUNETIX_URL,
+    BURP_URL,
+    ACUNETIX,
+    BURP,
+    SCANNER,
+    AcunetixSettings,
+    BurpSettings,
+    ScannerSettings,
+    get_acunetix_config,
+    get_burp_config,
+    print_config,
+)
+
+# Acunetix Integration
+from aipt_v2.tools.scanners.acunetix_tool import (
+    AcunetixTool,
+    AcunetixConfig,
+    ScanProfile,
+    ScanStatus as AcunetixScanStatus,
+    ScanResult as AcunetixScanResult,
+    Vulnerability as AcunetixVulnerability,
+    Severity as AcunetixSeverity,
+    get_acunetix,
+    acunetix_scan,
+    acunetix_status,
+    acunetix_vulns,
+    acunetix_summary,
+)
+
+# Burp Suite Integration
+from aipt_v2.tools.scanners.burp_tool import (
+    BurpTool,
+    BurpConfig,
+    ScanStatus as BurpScanStatus,
+    ScanResult as BurpScanResult,
+    Issue as BurpIssue,
+    IssueSeverity,
+    IssueConfidence,
+    get_burp,
+    burp_scan,
+    burp_status,
+    burp_issues,
+    burp_summary,
+)
+
+# Nessus Integration
+from aipt_v2.tools.scanners.nessus_tool import (
+    NessusTool,
+    NessusConfig,
+    ScanResult as NessusScanResult,
+    Vulnerability as NessusVulnerability,
+    get_nessus,
+    nessus_scan,
+    nessus_vulns,
+    nessus_summary,
+)
+
+# OWASP ZAP Integration
+from aipt_v2.tools.scanners.zap_tool import (
+    ZAPTool,
+    ZAPConfig,
+    ScanResult as ZAPScanResult,
+    Alert as ZAPAlert,
+    get_zap,
+    zap_scan,
+    zap_alerts,
+    zap_summary,
+)
+
+__all__ = [
+    # Configuration
+    "SCANNER_SERVER_IP",
+    "ACUNETIX_PORT",
+    "BURP_PORT",
+    "ACUNETIX_URL",
+    "BURP_URL",
+    "ACUNETIX",
+    "BURP",
+    "SCANNER",
+    "AcunetixSettings",
+    "BurpSettings",
+    "ScannerSettings",
+    "get_acunetix_config",
+    "get_burp_config",
+    "print_config",
+    # Acunetix
+    "AcunetixTool",
+    "AcunetixConfig",
+    "ScanProfile",
+    "AcunetixScanStatus",
+    "AcunetixScanResult",
+    "AcunetixVulnerability",
+    "AcunetixSeverity",
+    "get_acunetix",
+    "acunetix_scan",
+    "acunetix_status",
+    "acunetix_vulns",
+    "acunetix_summary",
+    # Burp Suite
+    "BurpTool",
+    "BurpConfig",
+    "BurpScanStatus",
+    "BurpScanResult",
+    "BurpIssue",
+    "IssueSeverity",
+    "IssueConfidence",
+    "get_burp",
+    "burp_scan",
+    "burp_status",
+    "burp_issues",
+    "burp_summary",
+    # Nessus
+    "NessusTool",
+    "NessusConfig",
+    "NessusScanResult",
+    "NessusVulnerability",
+    "get_nessus",
+    "nessus_scan",
+    "nessus_vulns",
+    "nessus_summary",
+    # OWASP ZAP
+    "ZAPTool",
+    "ZAPConfig",
+    "ZAPScanResult",
+    "ZAPAlert",
+    "get_zap",
+    "zap_scan",
+    "zap_alerts",
+    "zap_summary",
+    # Unified Interface
+    "ScannerType",
+    "get_scanner",
+    "scan_target",
+    "get_findings",
+    "get_all_findings",
+    "test_all_connections",
+]
+
+
+# ==================== Unified Scanner Interface ====================
+
+class ScannerType:
+    """Scanner type constants."""
+    ACUNETIX = "acunetix"
+    BURP = "burp"
+    NESSUS = "nessus"
+    ZAP = "zap"
+
+
+def get_scanner(scanner_type: str, config: dict = None):
+    """
+    Get a scanner instance by type.
+
+    Args:
+        scanner_type: 'acunetix', 'burp', 'nessus', or 'zap'
+        config: Optional configuration dict
+
+    Returns:
+        Scanner tool instance
+    """
+    if scanner_type == ScannerType.ACUNETIX:
+        cfg = AcunetixConfig(**config) if config else None
+        return get_acunetix(cfg)
+    elif scanner_type == ScannerType.BURP:
+        cfg = BurpConfig(**config) if config else None
+        return get_burp(cfg)
+    elif scanner_type == ScannerType.NESSUS:
+        cfg = NessusConfig(**config) if config else None
+        return get_nessus(cfg)
+    elif scanner_type == ScannerType.ZAP:
+        cfg = ZAPConfig(**config) if config else None
+        return get_zap(cfg)
+    else:
+        raise ValueError(f"Unknown scanner type: {scanner_type}")
+
+
+def scan_target(url: str, scanner_type: str = "acunetix", **kwargs):
+    """
+    Unified scan function - works with any scanner.
+
+    Args:
+        url: Target URL to scan
+        scanner_type: 'acunetix', 'burp', 'nessus', or 'zap'
+        **kwargs: Scanner-specific options
+
+    Returns:
+        ScanResult from the chosen scanner
+    """
+    if scanner_type == ScannerType.ACUNETIX:
+        profile = kwargs.get("profile", "full")
+        return acunetix_scan(url, profile)
+    elif scanner_type == ScannerType.BURP:
+        config_id = kwargs.get("config_id")
+        return burp_scan(url, config_id)
+    elif scanner_type == ScannerType.NESSUS:
+        policy = kwargs.get("policy", "basic")
+        return nessus_scan(url, policy)
+    elif scanner_type == ScannerType.ZAP:
+        return zap_scan(url)
+    else:
+        raise ValueError(f"Unknown scanner type: {scanner_type}")
+
+
+def get_findings(scanner_type: str, scan_id: str = None, severity: str = None):
+    """
+    Get findings from any scanner in unified AIPT format.
+
+    Args:
+        scanner_type: 'acunetix', 'burp', 'nessus', or 'zap'
+        scan_id: Optional scan ID filter
+        severity: Optional severity filter
+
+    Returns:
+        List of findings in AIPT format
+    """
+    if scanner_type == ScannerType.ACUNETIX:
+        return acunetix_vulns(scan_id, severity)
+    elif scanner_type == ScannerType.BURP:
+        return burp_issues(scan_id, severity)
+    elif scanner_type == ScannerType.NESSUS:
+        return nessus_vulns(scan_id, severity)
+    elif scanner_type == ScannerType.ZAP:
+        return zap_alerts(scan_id, severity)
+    else:
+        raise ValueError(f"Unknown scanner type: {scanner_type}")
+
+
+def get_all_findings(scan_ids: dict = None) -> list:
+    """
+    Aggregate findings from all scanners.
+
+    Args:
+        scan_ids: Dict mapping scanner_type to scan_id
+
+    Returns:
+        Combined list of findings from all scanners
+    """
+    findings = []
+    scan_ids = scan_ids or {}
+
+    try:
+        acunetix_findings = acunetix_vulns(scan_ids.get("acunetix"))
+        findings.extend(acunetix_findings)
+    except Exception:
+        pass  # Scanner not available
+
+    try:
+        burp_findings = burp_issues(scan_ids.get("burp"))
+        findings.extend(burp_findings)
+    except Exception:
+        pass  # Scanner not available
+
+    try:
+        nessus_findings = nessus_vulns(scan_ids.get("nessus"))
+        findings.extend(nessus_findings)
+    except Exception:
+        pass  # Scanner not available
+
+    try:
+        zap_findings = zap_alerts(scan_ids.get("zap"))
+        findings.extend(zap_findings)
+    except Exception:
+        pass  # Scanner not available
+
+    return findings
+
+
+def test_all_connections() -> dict:
+    """
+    Test connections to all configured scanners.
+
+    Returns:
+        Dict with connection status for each scanner
+    """
+    results = {}
+
+    # Test Acunetix
+    try:
+        acunetix = AcunetixTool()
+        results["acunetix"] = {
+            "connected": acunetix.connect(),
+            "url": acunetix.config.base_url,
+            "info": acunetix.get_info() if acunetix.is_connected() else None
+        }
+    except Exception as e:
+        results["acunetix"] = {"connected": False, "error": str(e)}
+
+    # Test Burp
+    try:
+        burp = BurpTool()
+        results["burp"] = {
+            "connected": burp.connect(),
+            "url": burp.config.base_url,
+            "info": burp.get_info() if burp.is_connected() else None
+        }
+    except Exception as e:
+        results["burp"] = {"connected": False, "error": str(e)}
+
+    # Test Nessus
+    try:
+        nessus = NessusTool()
+        results["nessus"] = {
+            "connected": nessus.connect(),
+            "url": nessus.config.base_url,
+            "info": nessus.get_info() if nessus.is_connected() else None
+        }
+    except Exception as e:
+        results["nessus"] = {"connected": False, "error": str(e)}
+
+    # Test OWASP ZAP
+    try:
+        zap = ZAPTool()
+        results["zap"] = {
+            "connected": zap.connect(),
+            "url": zap.config.base_url,
+            "info": zap.get_info() if zap.is_connected() else None
+        }
+    except Exception as e:
+        results["zap"] = {"connected": False, "error": str(e)}
+
+    return results