aiptx 2.0.7__py3-none-any.whl

aipt_v2/compliance/pci_mapping.py

@@ -0,0 +1,297 @@
+"""
+PCI-DSS 4.0 Mapping
+
+Payment Card Industry Data Security Standard mapping.
+Maps CWEs to PCI-DSS 4.0 requirements.
+
+Usage:
+    from aipt_v2.compliance import PCIMapper, get_pci_requirement
+
+    req = get_pci_requirement("CWE-89")  # Returns "6.2"
+"""
+
+from dataclasses import dataclass, field
+from typing import List, Dict, Optional
+
+
+@dataclass
+class PCIRequirement:
+    """PCI-DSS requirement definition."""
+    id: str
+    name: str
+    description: str
+    sub_requirements: List[str]
+    cwes: List[str]
+    testing_procedures: List[str]
+
+
+# PCI-DSS 4.0 Requirements (security-relevant subset)
+PCI_DSS_REQUIREMENTS = {
+    "2.2": PCIRequirement(
+        id="2.2",
+        name="System Components are Configured and Managed Securely",
+        description="Configuration standards are developed, implemented, and maintained "
+                   "for system components that are consistent with industry-accepted "
+                   "system hardening standards.",
+        sub_requirements=[
+            "2.2.1 - Configuration standards are implemented",
+            "2.2.2 - Vendor default accounts are managed",
+            "2.2.3 - Primary functions are separated on different servers",
+            "2.2.4 - Only necessary services and protocols are enabled",
+            "2.2.5 - Insecure services and protocols are addressed",
+            "2.2.6 - System security parameters are configured",
+            "2.2.7 - Non-console administrative access is encrypted"
+        ],
+        cwes=["CWE-16", "CWE-260", "CWE-611", "CWE-756", "CWE-942"],
+        testing_procedures=[
+            "Examine configuration standards",
+            "Examine system configurations",
+            "Interview system administrators"
+        ]
+    ),
+
+    "3.4": PCIRequirement(
+        id="3.4",
+        name="PAN is Rendered Unreadable Anywhere it is Stored",
+        description="Primary Account Numbers (PAN) is rendered unreadable anywhere "
+                   "it is stored using strong cryptography.",
+        sub_requirements=[
+            "3.4.1 - PAN rendered unreadable via cryptography or truncation",
+            "3.4.2 - Technical controls prevent PAN copy/relocation"
+        ],
+        cwes=["CWE-311", "CWE-312", "CWE-313", "CWE-316"],
+        testing_procedures=[
+            "Examine data repositories to verify PAN is unreadable",
+            "Examine cryptographic key management procedures"
+        ]
+    ),
+
+    "3.5": PCIRequirement(
+        id="3.5",
+        name="Cryptographic Keys are Protected",
+        description="Cryptographic keys used to protect stored account data "
+                   "are secured.",
+        sub_requirements=[
+            "3.5.1 - Key access restricted to fewest custodians",
+            "3.5.1.1 - Service providers: key per customer",
+            "3.5.1.2 - Key storage in secure cryptographic device",
+            "3.5.1.3 - Key stored in fewest locations"
+        ],
+        cwes=["CWE-320", "CWE-321", "CWE-326", "CWE-327"],
+        testing_procedures=[
+            "Examine documented procedures",
+            "Interview personnel",
+            "Observe key management processes"
+        ]
+    ),
+
+    "4.1": PCIRequirement(
+        id="4.1",
+        name="Strong Cryptography Protects Data During Transmission",
+        description="Processes and mechanisms are defined to protect account data "
+                   "with strong cryptography during transmission over open networks.",
+        sub_requirements=[
+            "4.1.1 - Security policies and procedures documented",
+            "4.1.2 - Roles and responsibilities assigned"
+        ],
+        cwes=["CWE-319", "CWE-326", "CWE-327", "CWE-523"],
+        testing_procedures=[
+            "Examine policies and configuration standards",
+            "Interview responsible personnel"
+        ]
+    ),
+
+    "6.2": PCIRequirement(
+        id="6.2",
+        name="Bespoke and Custom Software is Developed Securely",
+        description="Bespoke and custom software is developed securely, and "
+                   "software development follows secure software development guidance.",
+        sub_requirements=[
+            "6.2.1 - Software developed based on industry standards",
+            "6.2.2 - Personnel trained in secure development",
+            "6.2.3 - Code reviewed for vulnerabilities",
+            "6.2.4 - Common software attacks addressed"
+        ],
+        cwes=["CWE-79", "CWE-89", "CWE-78", "CWE-94", "CWE-502", "CWE-918", "CWE-22",
+              "CWE-434", "CWE-352", "CWE-601", "CWE-20"],
+        testing_procedures=[
+            "Examine secure development procedures",
+            "Interview developers",
+            "Review code review processes"
+        ]
+    ),
+
+    "6.3": PCIRequirement(
+        id="6.3",
+        name="Security Vulnerabilities are Identified and Addressed",
+        description="Security vulnerabilities are identified and addressed.",
+        sub_requirements=[
+            "6.3.1 - Vulnerabilities identified via reputable sources",
+            "6.3.2 - Inventory of bespoke software maintained",
+            "6.3.3 - Vulnerabilities addressed via patching"
+        ],
+        cwes=["CWE-937", "CWE-1035", "CWE-1104"],
+        testing_procedures=[
+            "Examine processes for identifying vulnerabilities",
+            "Interview responsible personnel"
+        ]
+    ),
+
+    "7.1": PCIRequirement(
+        id="7.1",
+        name="Access to System Components is Defined and Assigned",
+        description="Processes and mechanisms for restricting access to system "
+                   "components and cardholder data are defined and understood.",
+        sub_requirements=[
+            "7.1.1 - Security policies defined and known",
+            "7.1.2 - Access control model defined"
+        ],
+        cwes=["CWE-284", "CWE-285", "CWE-862", "CWE-863", "CWE-639"],
+        testing_procedures=[
+            "Examine documented policies",
+            "Interview personnel"
+        ]
+    ),
+
+    "8.3": PCIRequirement(
+        id="8.3",
+        name="Strong Authentication for Users and Administrators",
+        description="Strong authentication for users and administrators is established "
+                   "and managed.",
+        sub_requirements=[
+            "8.3.1 - All access authenticated",
+            "8.3.2 - Strong cryptography for authentication",
+            "8.3.4 - MFA for remote network access",
+            "8.3.5 - MFA for all access to CDE",
+            "8.3.6 - MFA systems cannot be bypassed",
+            "8.3.9 - Passwords have minimum complexity",
+            "8.3.10 - Failed login attempts limited"
+        ],
+        cwes=["CWE-287", "CWE-521", "CWE-798", "CWE-307", "CWE-384", "CWE-306"],
+        testing_procedures=[
+            "Examine system configuration standards",
+            "Observe authentication processes"
+        ]
+    ),
+
+    "10.2": PCIRequirement(
+        id="10.2",
+        name="Audit Logs are Implemented",
+        description="Audit logs are implemented to support the detection of "
+                   "anomalies and suspicious activity.",
+        sub_requirements=[
+            "10.2.1 - Audit logs enabled and active",
+            "10.2.1.1 - User access to CHD logged",
+            "10.2.1.2 - Admin actions logged",
+            "10.2.1.3 - Access to audit logs logged",
+            "10.2.1.4 - Invalid access attempts logged",
+            "10.2.1.5 - Changes to auth credentials logged",
+            "10.2.1.6 - System/log events logged",
+            "10.2.1.7 - Security events logged"
+        ],
+        cwes=["CWE-778", "CWE-223", "CWE-117", "CWE-532"],
+        testing_procedures=[
+            "Examine audit log configurations",
+            "Interview personnel",
+            "Review log samples"
+        ]
+    ),
+
+    "11.3": PCIRequirement(
+        id="11.3",
+        name="External and Internal Vulnerabilities Regularly Identified",
+        description="External and internal vulnerabilities are regularly identified, "
+                   "prioritized, and addressed.",
+        sub_requirements=[
+            "11.3.1 - Internal vulnerability scans quarterly",
+            "11.3.2 - External vulnerability scans quarterly",
+            "11.3.3 - Vulnerability scans after significant changes",
+            "11.3.4 - Internal penetration testing"
+        ],
+        cwes=["CWE-937", "CWE-1104"],
+        testing_procedures=[
+            "Examine scan reports",
+            "Interview responsible personnel",
+            "Examine remediation processes"
+        ]
+    )
+}
+
+
+class PCIMapper:
+    """PCI-DSS specific mapper."""
+
+    def __init__(self):
+        self.requirements = PCI_DSS_REQUIREMENTS
+
+    def get_requirement(self, cwe_id: str) -> Optional[PCIRequirement]:
+        """Get PCI requirement for a CWE."""
+        cwe_id = cwe_id.upper()
+        if not cwe_id.startswith("CWE-"):
+            cwe_id = f"CWE-{cwe_id}"
+
+        for req_id, requirement in self.requirements.items():
+            if cwe_id in requirement.cwes:
+                return requirement
+
+        return None
+
+    def get_requirement_by_id(self, req_id: str) -> Optional[PCIRequirement]:
+        """Get PCI requirement by ID."""
+        return self.requirements.get(req_id)
+
+    def get_all_requirements_for_cwe(self, cwe_id: str) -> List[PCIRequirement]:
+        """Get all PCI requirements mapped to a CWE."""
+        cwe_id = cwe_id.upper()
+        if not cwe_id.startswith("CWE-"):
+            cwe_id = f"CWE-{cwe_id}"
+
+        requirements = []
+        for requirement in self.requirements.values():
+            if cwe_id in requirement.cwes:
+                requirements.append(requirement)
+
+        return requirements
+
+    def get_compliance_status(self, findings: List[Dict]) -> Dict[str, Dict]:
+        """
+        Calculate PCI-DSS compliance status based on findings.
+
+        Args:
+            findings: List of findings with CWE IDs
+
+        Returns:
+            Dict with compliance status per requirement
+        """
+        status = {}
+
+        for req_id, requirement in self.requirements.items():
+            affected_cwes = []
+            for finding in findings:
+                cwe = finding.get("cwe", finding.get("cwe_id", ""))
+                if cwe.upper() in requirement.cwes or f"CWE-{cwe}" in requirement.cwes:
+                    affected_cwes.append(cwe)
+
+            status[req_id] = {
+                "requirement_name": requirement.name,
+                "compliant": len(affected_cwes) == 0,
+                "findings_count": len(affected_cwes),
+                "affected_cwes": affected_cwes
+            }
+
+        return status
+
+
+def get_pci_requirement(cwe_id: str) -> Optional[str]:
+    """
+    Get PCI-DSS requirement ID for a CWE.
+
+    Args:
+        cwe_id: CWE identifier
+
+    Returns:
+        PCI requirement ID or None
+    """
+    mapper = PCIMapper()
+    requirement = mapper.get_requirement(cwe_id)
+    return requirement.id if requirement else None

aipt_v2/config.py

@@ -0,0 +1,341 @@
+"""
+AIPT v2 Configuration Management
+================================
+
+Centralized configuration with validation using Pydantic.
+Loads from environment variables with sensible defaults.
+Also loads from ~/.aiptx/.env file created by the setup wizard.
+"""
+
+import os
+from typing import List, Optional
+from functools import lru_cache
+from pathlib import Path
+
+from dotenv import load_dotenv
+from pydantic import BaseModel, Field, field_validator, model_validator
+from pydantic_settings import BaseSettings
+
+from .utils.logging import logger
+
+
+def _load_config_files():
+    """
+    Load configuration from .env files.
+
+    Priority (highest to lowest):
+    1. Environment variables (already set in shell)
+    2. Local .env file in current directory
+    3. Global ~/.aiptx/.env file from setup wizard
+    """
+    # Load global config first (lowest priority)
+    global_env = Path.home() / ".aiptx" / ".env"
+    if global_env.exists():
+        load_dotenv(global_env, override=False)
+
+    # Load local .env file (higher priority, but doesn't override existing env vars)
+    local_env = Path(".env")
+    if local_env.exists():
+        load_dotenv(local_env, override=False)
+
+
+# Load config files on module import
+_load_config_files()
+
+
+class LLMSettings(BaseModel):
+    """LLM provider configuration."""
+
+    provider: str = Field(default="anthropic", description="LLM provider name")
+    model: str = Field(default="claude-sonnet-4-20250514", description="Model identifier")
+    api_key: Optional[str] = Field(default=None, description="API key")
+    api_base: Optional[str] = Field(default=None, description="Custom API base URL")
+    timeout: int = Field(default=120, ge=10, le=600, description="Request timeout in seconds")
+    max_tokens: int = Field(default=4096, ge=100, le=128000, description="Max response tokens")
+    temperature: float = Field(default=0.7, ge=0.0, le=2.0, description="Sampling temperature")
+    enable_caching: bool = Field(default=True, description="Enable prompt caching")
+
+    @field_validator("api_key", mode="before")
+    @classmethod
+    def get_api_key_from_env(cls, v):
+        if v:
+            return v
+        # Try common environment variables
+        for key in ["ANTHROPIC_API_KEY", "OPENAI_API_KEY", "DEEPSEEK_API_KEY", "LLM_API_KEY"]:
+            if os.getenv(key):
+                return os.getenv(key)
+        return None
+
+
+class ScannerSettings(BaseModel):
+    """Enterprise scanner configuration."""
+
+    # Acunetix
+    acunetix_url: Optional[str] = Field(default=None, description="Acunetix API URL")
+    acunetix_api_key: Optional[str] = Field(default=None, description="Acunetix API key")
+
+    # Burp Suite
+    burp_url: Optional[str] = Field(default=None, description="Burp Suite API URL")
+    burp_api_key: Optional[str] = Field(default=None, description="Burp Suite API key")
+
+    # Nessus
+    nessus_url: Optional[str] = Field(default=None, description="Nessus API URL")
+    nessus_access_key: Optional[str] = Field(default=None, description="Nessus access key")
+    nessus_secret_key: Optional[str] = Field(default=None, description="Nessus secret key")
+
+    # OWASP ZAP
+    zap_url: Optional[str] = Field(default=None, description="ZAP API URL")
+    zap_api_key: Optional[str] = Field(default=None, description="ZAP API key")
+
+    @field_validator("acunetix_url", "burp_url", "nessus_url", "zap_url", mode="before")
+    @classmethod
+    def validate_url(cls, v):
+        if v and not v.startswith(("http://", "https://")):
+            # Auto-prepend http:// if no scheme provided (user-friendly)
+            v = f"http://{v}"
+        return v
+
+
+class VPSSettings(BaseModel):
+    """VPS execution configuration."""
+
+    host: Optional[str] = Field(default=None, description="VPS hostname or IP")
+    user: str = Field(default="ubuntu", description="SSH username")
+    key_path: Optional[str] = Field(default=None, description="Path to SSH private key")
+    port: int = Field(default=22, ge=1, le=65535, description="SSH port")
+    # Security: Use /var/tmp for longer-lived temp files (survives reboots)
+    # The directory should be created with restricted permissions (700) on the VPS
+    results_dir: str = Field(
+        default="/var/tmp/aipt_results",
+        description="Remote results directory (should be created with mode 700)"
+    )
+    timeout: int = Field(default=300, ge=30, le=3600, description="Command timeout in seconds")
+
+    @field_validator("key_path", mode="before")
+    @classmethod
+    def expand_key_path(cls, v):
+        if v:
+            return str(Path(v).expanduser().resolve())
+        return v
+
+
+class APISettings(BaseModel):
+    """REST API configuration."""
+
+    # Security: Default to localhost to avoid accidental exposure (CWE-605)
+    # Use AIPT_API__HOST=0.0.0.0 for production deployments behind a reverse proxy
+    host: str = Field(default="127.0.0.1", description="API host (use 0.0.0.0 for network access)")
+    port: int = Field(default=8000, ge=1, le=65535, description="API port")
+    cors_origins: List[str] = Field(
+        default=["http://localhost:3000"],
+        description="Allowed CORS origins"
+    )
+    rate_limit: str = Field(default="100/minute", description="Rate limit per client")
+    enable_docs: bool = Field(default=True, description="Enable Swagger/OpenAPI docs")
+
+
+class DatabaseSettings(BaseModel):
+    """Database configuration."""
+
+    url: str = Field(default="sqlite:///./aipt.db", description="Database connection URL")
+    echo: bool = Field(default=False, description="Echo SQL queries")
+    pool_size: int = Field(default=5, ge=1, le=100, description="Connection pool size")
+
+
+class LoggingSettings(BaseModel):
+    """Logging configuration."""
+
+    level: str = Field(default="INFO", description="Log level")
+    format: str = Field(default="console", description="Log format (console or json)")
+    redact_secrets: bool = Field(default=True, description="Redact sensitive values in logs")
+
+
+class AIPTConfig(BaseSettings):
+    """
+    Main AIPT v2 configuration.
+
+    Loads from environment variables with AIPT_ prefix.
+    """
+
+    # Sub-configurations
+    llm: LLMSettings = Field(default_factory=LLMSettings)
+    scanners: ScannerSettings = Field(default_factory=ScannerSettings)
+    vps: VPSSettings = Field(default_factory=VPSSettings)
+    api: APISettings = Field(default_factory=APISettings)
+    database: DatabaseSettings = Field(default_factory=DatabaseSettings)
+    logging: LoggingSettings = Field(default_factory=LoggingSettings)
+
+    # General settings
+    sandbox_mode: bool = Field(default=False, description="Run tools in sandbox")
+    output_dir: Path = Field(default=Path("./results"), description="Output directory")
+    reports_dir: Path = Field(default=Path("./reports"), description="Reports directory")
+
+    model_config = {
+        "env_prefix": "AIPT_",
+        "env_nested_delimiter": "__",
+        "case_sensitive": False,
+    }
+
+    @model_validator(mode="after")
+    def create_directories(self):
+        """Ensure output directories exist."""
+        self.output_dir.mkdir(parents=True, exist_ok=True)
+        self.reports_dir.mkdir(parents=True, exist_ok=True)
+        return self
+
+
+@lru_cache(maxsize=1)
+def get_config() -> AIPTConfig:
+    """
+    Get the global configuration instance.
+
+    Returns:
+        AIPTConfig instance loaded from environment
+    """
+    # Load from environment
+    # Note: Setup wizard saves with AIPT_ prefix and __ delimiter (e.g., AIPT_LLM__PROVIDER)
+    # We check both formats for backwards compatibility
+    config = AIPTConfig(
+        llm=LLMSettings(
+            provider=os.getenv("AIPT_LLM__PROVIDER") or os.getenv("AIPT_LLM_PROVIDER", "anthropic"),
+            model=os.getenv("AIPT_LLM__MODEL") or os.getenv("AIPT_LLM_MODEL", "claude-sonnet-4-20250514"),
+            api_key=os.getenv("ANTHROPIC_API_KEY") or os.getenv("OPENAI_API_KEY") or os.getenv("DEEPSEEK_API_KEY") or os.getenv("LLM_API_KEY"),
+            timeout=int(os.getenv("AIPT_LLM__TIMEOUT") or os.getenv("AIPT_LLM_TIMEOUT", "120")),
+        ),
+        scanners=ScannerSettings(
+            acunetix_url=os.getenv("AIPT_SCANNERS__ACUNETIX_URL") or os.getenv("ACUNETIX_URL"),
+            acunetix_api_key=os.getenv("AIPT_SCANNERS__ACUNETIX_API_KEY") or os.getenv("ACUNETIX_API_KEY"),
+            burp_url=os.getenv("AIPT_SCANNERS__BURP_URL") or os.getenv("BURP_URL"),
+            burp_api_key=os.getenv("AIPT_SCANNERS__BURP_API_KEY") or os.getenv("BURP_API_KEY"),
+            nessus_url=os.getenv("AIPT_SCANNERS__NESSUS_URL") or os.getenv("NESSUS_URL"),
+            nessus_access_key=os.getenv("AIPT_SCANNERS__NESSUS_ACCESS_KEY") or os.getenv("NESSUS_ACCESS_KEY"),
+            nessus_secret_key=os.getenv("AIPT_SCANNERS__NESSUS_SECRET_KEY") or os.getenv("NESSUS_SECRET_KEY"),
+            zap_url=os.getenv("AIPT_SCANNERS__ZAP_URL") or os.getenv("ZAP_URL"),
+            zap_api_key=os.getenv("AIPT_SCANNERS__ZAP_API_KEY") or os.getenv("ZAP_API_KEY"),
+        ),
+        vps=VPSSettings(
+            host=os.getenv("AIPT_VPS__HOST") or os.getenv("VPS_HOST"),
+            user=os.getenv("AIPT_VPS__USER") or os.getenv("VPS_USER", "ubuntu"),
+            key_path=os.getenv("AIPT_VPS__KEY_PATH") or os.getenv("VPS_KEY"),
+        ),
+        api=APISettings(
+            cors_origins=os.getenv("AIPT_CORS_ORIGINS", "http://localhost:3000").split(","),
+            rate_limit=os.getenv("AIPT_RATE_LIMIT", "100/minute"),
+        ),
+        database=DatabaseSettings(
+            url=os.getenv("DATABASE_URL", "sqlite:///./aipt.db"),
+        ),
+        logging=LoggingSettings(
+            level=os.getenv("AIPT_LOG_LEVEL", "WARNING"),  # Default to WARNING for cleaner output
+            format=os.getenv("AIPT_LOG_FORMAT", "console"),
+        ),
+        sandbox_mode=os.getenv("AIPT_SANDBOX_MODE", "false").lower() == "true",
+    )
+
+    # Only log configuration details at DEBUG level for cleaner default output
+    logger.debug(
+        "Configuration loaded",
+        llm_provider=config.llm.provider,
+        llm_model=config.llm.model,
+        has_llm_key=bool(config.llm.api_key),
+        has_acunetix=bool(config.scanners.acunetix_url),
+        has_burp=bool(config.scanners.burp_url),
+        has_nessus=bool(config.scanners.nessus_url),
+        has_vps=bool(config.vps.host),
+        sandbox_mode=config.sandbox_mode,
+    )
+
+    return config
+
+
+def reload_config() -> AIPTConfig:
+    """
+    Reload configuration from files and environment.
+
+    This clears the cached config and reloads from:
+    - ~/.aiptx/.env (setup wizard config)
+    - ./.env (local project config)
+    - Environment variables
+
+    Useful after running the setup wizard.
+
+    Returns:
+        Fresh AIPTConfig instance
+    """
+    # Clear the cached config
+    get_config.cache_clear()
+
+    # Reload .env files
+    _load_config_files()
+
+    # Return fresh config
+    return get_config()
+
+
+def validate_config_for_features(features: List[str]) -> List[str]:
+    """
+    Validate that required configuration is present for requested features.
+
+    Args:
+        features: List of feature names to validate
+
+    Returns:
+        List of error messages (empty if all valid)
+    """
+    config = get_config()
+    errors = []
+
+    if "llm" in features and not config.llm.api_key:
+        errors.append(
+            "LLM API key required. Set ANTHROPIC_API_KEY or OPENAI_API_KEY environment variable."
+        )
+
+    if "acunetix" in features:
+        if not config.scanners.acunetix_url:
+            errors.append("Acunetix URL required. Set ACUNETIX_URL environment variable.")
+        if not config.scanners.acunetix_api_key:
+            errors.append("Acunetix API key required. Set ACUNETIX_API_KEY environment variable.")
+
+    if "burp" in features:
+        if not config.scanners.burp_url:
+            errors.append("Burp Suite URL required. Set BURP_URL environment variable.")
+        if not config.scanners.burp_api_key:
+            errors.append("Burp Suite API key required. Set BURP_API_KEY environment variable.")
+
+    if "nessus" in features:
+        if not config.scanners.nessus_url:
+            errors.append("Nessus URL required. Set NESSUS_URL environment variable.")
+        if not config.scanners.nessus_access_key:
+            errors.append("Nessus access key required. Set NESSUS_ACCESS_KEY environment variable.")
+        if not config.scanners.nessus_secret_key:
+            errors.append("Nessus secret key required. Set NESSUS_SECRET_KEY environment variable.")
+
+    if "vps" in features:
+        if not config.vps.host:
+            errors.append("VPS host required. Set VPS_HOST environment variable.")
+        if not config.vps.key_path:
+            errors.append("VPS SSH key path required. Set VPS_KEY environment variable.")
+        elif not Path(config.vps.key_path).exists():
+            errors.append(f"VPS SSH key not found: {config.vps.key_path}")
+
+    return errors
+
+
+# Convenience function
+def require_config(*features: str) -> AIPTConfig:
+    """
+    Get config and raise if required features are not configured.
+
+    Args:
+        *features: Feature names to validate
+
+    Returns:
+        AIPTConfig instance
+
+    Raises:
+        ValueError: If required configuration is missing
+    """
+    errors = validate_config_for_features(list(features))
+    if errors:
+        raise ValueError("Configuration errors:\n" + "\n".join(f"  - {e}" for e in errors))
+    return get_config()

aipt_v2/core/__init__.py

@@ -0,0 +1,43 @@
+"""
+AIPT Core Module
+
+LangGraph-based autonomous pentesting agent with:
+- Multi-provider LLM abstraction (OpenAI, Anthropic, Ollama)
+- Memory management with automatic compression
+- State machine for pentest workflow
+
+Inspired by: Strix (LangGraph state machine, 300 iterations)
+"""
+from __future__ import annotations
+
+from .llm import (
+    LLMProvider,
+    LLMResponse,
+    OpenAIProvider,
+    AnthropicProvider,
+    OllamaProvider,
+    get_llm,
+)
+from .memory import MemoryManager, MemoryConfig
+from .agent import AIPTAgent, PentestState, Phase
+from .ptt import PTTTracker, PTTNode
+
+__all__ = [
+    # LLM
+    "LLMProvider",
+    "LLMResponse",
+    "OpenAIProvider",
+    "AnthropicProvider",
+    "OllamaProvider",
+    "get_llm",
+    # Memory
+    "MemoryManager",
+    "MemoryConfig",
+    # Agent
+    "AIPTAgent",
+    "PentestState",
+    "Phase",
+    # PTT
+    "PTTTracker",
+    "PTTNode",
+]