aiptx 2.0.2__py3-none-any.whl

aipt_v2/__init__.py

@@ -0,0 +1,110 @@
+"""
+AIPT v2 - AI-Powered Penetration Testing Framework
+===================================================
+
+A unified penetration testing framework built on top of 8 reference tools:
+- AIPTx: LLM (litellm), Runtime (Docker), Tools (Browser, Terminal, Proxy)
+- pentest-agent: CVE Intelligence with EPSS scoring
+- PentestAssistant: RAG-based tool selection with BGE embeddings
+- PentestGPT: PTT (Penetration Testing Tree) task tracking
+- VulnBot: Output parsing patterns
+- HackSynth: Multi-step reasoning
+- Pentagi: Docker isolation
+- ez-ai-agent: Simple execution model
+
+Features:
+- Universal LLM support via litellm (100+ models)
+- Docker sandbox execution
+- Browser automation via Playwright
+- Proxy interception via mitmproxy
+- CVE prioritization (CVSS + EPSS + trending + POC)
+- RAG tool selection with semantic search
+- Hierarchical task tracking
+- SQLAlchemy persistence
+- FastAPI REST API
+"""
+
+__version__ = "2.0.2"
+__author__ = "AIPT Team"
+
+# Available submodules (direct import)
+__all__ = [
+    # Core - LangGraph agent, LLM providers, memory
+    "core",
+    # Docker - Container management and sandboxing
+    "docker",
+    # Execution - Terminal, parser, sandbox integration
+    "execution",
+    # Orchestration - Pipeline, scheduler, progress tracking
+    "orchestration",
+    # Intelligence - Vulnerability analysis, triage, scope
+    "intelligence",
+    # Tools - Scanner integrations (Acunetix, Burp, etc.)
+    "tools",
+    # Payloads - XSS, SQLi, SSRF, SSTI, etc.
+    "payloads",
+    # Scanners - Nuclei, Nmap, Nikto wrappers
+    "scanners",
+    # Recon - Subdomain, DNS, tech detection
+    "recon",
+    # Browser - Playwright automation
+    "browser",
+    # Terminal - Command execution
+    "terminal",
+    # Proxy - mitmproxy interception
+    "proxy",
+]
+
+# Lazy imports to avoid failures when optional dependencies are missing
+
+
+def __getattr__(name):
+    """Lazy import handler for optional dependencies"""
+    if name == "LLM":
+        from aipt_v2.llm.llm import LLM
+        return LLM
+    elif name == "LLMConfig":
+        from aipt_v2.llm.config import LLMConfig
+        return LLMConfig
+    elif name == "PTT":
+        from aipt_v2.agents.ptt import PTT
+        return PTT
+    elif name == "BaseAgent":
+        from aipt_v2.agents.base import BaseAgent
+        return BaseAgent
+    elif name == "CVEIntelligence":
+        from aipt_v2.intelligence.cve_aipt import CVEIntelligence
+        return CVEIntelligence
+    elif name == "ToolRAG":
+        from aipt_v2.intelligence.rag import ToolRAG
+        return ToolRAG
+    elif name == "OutputParser":
+        from aipt_v2.tools.parser import OutputParser
+        return OutputParser
+    elif name == "Repository":
+        from aipt_v2.database.repository import Repository
+        return Repository
+    # New models module
+    elif name == "Finding":
+        from aipt_v2.models.findings import Finding
+        return Finding
+    elif name == "Severity":
+        from aipt_v2.models.findings import Severity
+        return Severity
+    elif name == "ScanConfig":
+        from aipt_v2.models.scan_config import ScanConfig
+        return ScanConfig
+    elif name == "ScanMode":
+        from aipt_v2.models.scan_config import ScanMode
+        return ScanMode
+    elif name == "PhaseResult":
+        from aipt_v2.models.phase_result import PhaseResult
+        return PhaseResult
+    # Reports module
+    elif name == "ReportGenerator":
+        from aipt_v2.reports.generator import ReportGenerator
+        return ReportGenerator
+    elif name == "ReportConfig":
+        from aipt_v2.reports.generator import ReportConfig
+        return ReportConfig
+    raise AttributeError(f"module 'aipt_v2' has no attribute '{name}'")

aipt_v2/__main__.py

@@ -0,0 +1,24 @@
+#!/usr/bin/env python3
+"""
+AIPTX - AI-Powered Penetration Testing Framework
+================================================
+
+This module allows running aiptx as a module:
+    python -m aiptx scan example.com
+    python -m aiptx --help
+
+Or directly after pipx install:
+    aiptx scan example.com
+"""
+
+import sys
+
+
+def main():
+    """Entry point for module execution."""
+    from cli import main as cli_main
+    sys.exit(cli_main())
+
+
+if __name__ == "__main__":
+    main()

aipt_v2/agents/AIPTxAgent/__init__.py

@@ -0,0 +1,10 @@
+"""
+AIPT AIPTxAgent - Main penetration testing agent
+"""
+
+from aipt_v2.agents.AIPTxAgent.aiptx_agent import AIPTxAgent
+
+# Backwards compatibility alias
+StrixAgent = AIPTxAgent
+
+__all__ = ["AIPTxAgent", "StrixAgent"]

aipt_v2/agents/AIPTxAgent/aiptx_agent.py

@@ -0,0 +1,211 @@
+"""
+AIPT AIPTxAgent - Main penetration testing agent
+
+This is the primary agent that orchestrates penetration testing activities.
+It uses the BaseAgent infrastructure with security-focused tools and prompts.
+"""
+
+import asyncio
+import logging
+from typing import Any, Optional, Dict
+
+from aipt_v2.agents.base import BaseAgent
+from aipt_v2.agents.ptt import PTT, TaskStatus
+from aipt_v2.llm.config import LLMConfig
+
+
+logger = logging.getLogger(__name__)
+
+
+class AIPTxAgent(BaseAgent):
+    """
+    AIPTxAgent - AI-powered penetration testing agent.
+
+    This agent performs autonomous security testing using a Think-Select-Execute-Learn loop:
+    1. THINK: Analyze current state and decide next action
+    2. SELECT: Choose appropriate security tools via RAG
+    3. EXECUTE: Run tools and capture output
+    4. LEARN: Extract findings, update PTT, decide next phase
+
+    Features:
+    - Multi-phase pentest tracking (recon, enum, exploit, post, report)
+    - RAG-based tool selection with 50+ security tools
+    - CVE intelligence with CVSS+EPSS+POC scoring
+    - Docker sandbox for isolated tool execution
+    - Browser automation for web application testing
+    - Proxy interception for traffic analysis
+    """
+
+    agent_name = "AIPTxAgent"
+    max_iterations = 300
+
+    def __init__(self, config: Dict[str, Any]):
+        """
+        Initialize AIPTxAgent.
+
+        Args:
+            config: Agent configuration with:
+                - llm_config: LLMConfig instance
+                - max_iterations: Maximum agent loop iterations
+                - non_interactive: Run without user interaction
+                - local_sources: Local source directories to mount
+        """
+        # Ensure llm_config is provided
+        if "llm_config" not in config:
+            config["llm_config"] = LLMConfig()
+
+        super().__init__(config)
+
+        # Initialize PTT for tracking pentest progress
+        self.ptt = PTT()
+
+        # Store scan configuration
+        self.scan_config: Optional[Dict[str, Any]] = None
+        self.targets_info: list[Dict[str, Any]] = []
+
+        # Results storage
+        self.findings: list[Dict[str, Any]] = []
+        self.vulnerabilities: list[Dict[str, Any]] = []
+
+    async def execute_scan(self, scan_config: Dict[str, Any]) -> Dict[str, Any]:
+        """
+        Execute a penetration test scan.
+
+        Args:
+            scan_config: Scan configuration with:
+                - scan_id: Unique scan identifier
+                - targets: List of target info dicts
+                - user_instructions: Optional user instructions
+                - run_name: Name for this run
+
+        Returns:
+            Dict with scan results including findings and vulnerabilities
+        """
+        self.scan_config = scan_config
+        self.targets_info = scan_config.get("targets", [])
+
+        # Build the task prompt
+        task = self._build_task_prompt(scan_config)
+
+        # Initialize PTT for first target
+        if self.targets_info:
+            first_target = self.targets_info[0].get("original", "unknown")
+            self.ptt.initialize(first_target)
+
+        logger.info(f"Starting penetration test scan: {scan_config.get('scan_id', 'unknown')}")
+
+        try:
+            # Run the agent loop
+            result = await self.agent_loop(task)
+
+            # Compile final results
+            final_result = {
+                "success": True,
+                "scan_id": scan_config.get("scan_id"),
+                "findings": self.findings,
+                "vulnerabilities": self.vulnerabilities,
+                "ptt_summary": self.ptt.get_summary() if self.ptt.target else {},
+                "agent_summary": self.state.get_execution_summary(),
+            }
+
+            return final_result
+
+        except asyncio.CancelledError:
+            logger.warning("Scan was cancelled")
+            return {
+                "success": False,
+                "error": "Scan was cancelled",
+                "scan_id": scan_config.get("scan_id"),
+            }
+        except Exception as e:
+            logger.exception(f"Scan failed: {e}")
+            return {
+                "success": False,
+                "error": str(e),
+                "scan_id": scan_config.get("scan_id"),
+            }
+
+    def _build_task_prompt(self, scan_config: Dict[str, Any]) -> str:
+        """Build the initial task prompt for the agent."""
+        targets = scan_config.get("targets", [])
+        user_instructions = scan_config.get("user_instructions", "")
+
+        # Build target description
+        if len(targets) == 1:
+            target_desc = targets[0].get("original", "unknown target")
+            target_type = targets[0].get("type", "unknown")
+        else:
+            target_desc = f"{len(targets)} targets"
+            target_type = "multiple"
+
+        task = f"""You are an AI penetration testing agent. Your mission is to perform a comprehensive security assessment on: {target_desc}
+
+Target Type: {target_type}
+"""
+
+        if len(targets) == 1:
+            details = targets[0].get("details", {})
+            if details:
+                task += f"Target Details: {details}\n"
+
+        if user_instructions:
+            task += f"\nUser Instructions: {user_instructions}\n"
+
+        task += """
+Your objectives:
+1. RECONNAISSANCE: Gather information about the target (ports, services, technologies)
+2. ENUMERATION: Identify potential attack vectors and vulnerabilities
+3. EXPLOITATION: Safely test identified vulnerabilities (do not cause damage)
+4. DOCUMENTATION: Record all findings with severity and remediation advice
+
+Guidelines:
+- Follow responsible disclosure practices
+- Document all findings clearly
+- Prioritize high-impact vulnerabilities
+- Stay within authorized scope
+- Use appropriate tools for each phase
+
+Begin your security assessment now. Start with reconnaissance to understand the target.
+"""
+
+        return task
+
+    def add_finding(self, finding: Dict[str, Any]) -> None:
+        """Add a finding to the scan results."""
+        self.findings.append(finding)
+
+        # Also track in PTT
+        if self.ptt.target:
+            phase = finding.get("phase", self.ptt.current_phase)
+            self.ptt.add_findings(phase, [finding])
+
+    def add_vulnerability(self, vulnerability: Dict[str, Any]) -> None:
+        """Add a vulnerability to the scan results."""
+        self.vulnerabilities.append(vulnerability)
+
+        # Also add as finding
+        self.add_finding({
+            **vulnerability,
+            "type": "vulnerability",
+        })
+
+        # Notify tracer if available
+        from aipt_v2.telemetry.tracer import get_global_tracer
+        tracer = get_global_tracer()
+        if tracer and hasattr(tracer, "report_vulnerability"):
+            tracer.report_vulnerability(
+                report_id=vulnerability.get("id", "VULN"),
+                title=vulnerability.get("title", "Unknown Vulnerability"),
+                content=vulnerability.get("description", ""),
+                severity=vulnerability.get("severity", "info"),
+            )
+
+    def get_ptt_summary(self) -> str:
+        """Get PTT progress summary for the LLM."""
+        if self.ptt.target:
+            return self.ptt.to_prompt()
+        return "No PTT initialized"
+
+
+# Backwards compatibility alias
+StrixAgent = AIPTxAgent

aipt_v2/agents/__init__.py

@@ -0,0 +1,24 @@
+"""
+AIPT Agents Module - Agent orchestration and task tracking
+"""
+
+# Core components that don't require external dependencies
+from aipt_v2.agents.ptt import PTT, Task, Phase, TaskStatus, PhaseType
+from aipt_v2.agents.state import AgentState
+
+__all__ = [
+    "PTT",
+    "Task",
+    "Phase",
+    "PhaseType",
+    "TaskStatus",
+    "AgentState",
+]
+
+
+def __getattr__(name):
+    """Lazy import for components with external dependencies"""
+    if name == "BaseAgent":
+        from aipt_v2.agents.base import BaseAgent
+        return BaseAgent
+    raise AttributeError(f"module 'aipt_v2.agents' has no attribute '{name}'")