aiptx 2.0.2__py3-none-any.whl

aipt_v2/llm/memory.py

@@ -0,0 +1,214 @@
+from __future__ import annotations
+
+import logging
+import os
+from typing import Any
+
+import litellm
+
+
+logger = logging.getLogger(__name__)
+
+
+MAX_TOTAL_TOKENS = 100_000
+MIN_RECENT_MESSAGES = 15
+
+SUMMARY_PROMPT_TEMPLATE = """You are an agent performing context
+condensation for a security agent. Your job is to compress scan data while preserving
+ALL operationally critical information for continuing the security assessment.
+
+CRITICAL ELEMENTS TO PRESERVE:
+- Discovered vulnerabilities and potential attack vectors
+- Scan results and tool outputs (compressed but maintaining key findings)
+- Access credentials, tokens, or authentication details found
+- System architecture insights and potential weak points
+- Progress made in the assessment
+- Failed attempts and dead ends (to avoid duplication)
+- Any decisions made about the testing approach
+
+COMPRESSION GUIDELINES:
+- Preserve exact technical details (URLs, paths, parameters, payloads)
+- Summarize verbose tool outputs while keeping critical findings
+- Maintain version numbers, specific technologies identified
+- Keep exact error messages that might indicate vulnerabilities
+- Compress repetitive or similar findings into consolidated form
+
+Remember: Another security agent will use this summary to continue the assessment.
+They must be able to pick up exactly where you left off without losing any
+operational advantage or context needed to find vulnerabilities.
+
+CONVERSATION SEGMENT TO SUMMARIZE:
+{conversation}
+
+Provide a technically precise summary that preserves all operational security context while
+keeping the summary concise and to the point."""
+
+
+def _count_tokens(text: str, model: str) -> int:
+    try:
+        count = litellm.token_counter(model=model, text=text)
+        return int(count)
+    except Exception:
+        logger.exception("Failed to count tokens")
+        return len(text) // 4  # Rough estimate
+
+
+def _get_message_tokens(msg: dict[str, Any], model: str) -> int:
+    content = msg.get("content", "")
+    if isinstance(content, str):
+        return _count_tokens(content, model)
+    if isinstance(content, list):
+        return sum(
+            _count_tokens(item.get("text", ""), model)
+            for item in content
+            if isinstance(item, dict) and item.get("type") == "text"
+        )
+    return 0
+
+
+def _extract_message_text(msg: dict[str, Any]) -> str:
+    content = msg.get("content", "")
+    if isinstance(content, str):
+        return content
+
+    if isinstance(content, list):
+        parts = []
+        for item in content:
+            if isinstance(item, dict):
+                if item.get("type") == "text":
+                    parts.append(item.get("text", ""))
+                elif item.get("type") == "image_url":
+                    parts.append("[IMAGE]")
+        return " ".join(parts)
+
+    return str(content)
+
+
+def _summarize_messages(
+    messages: list[dict[str, Any]],
+    model: str,
+    timeout: int = 600,
+) -> dict[str, Any]:
+    if not messages:
+        empty_summary = "<context_summary message_count='0'>{text}</context_summary>"
+        return {
+            "role": "assistant",
+            "content": empty_summary.format(text="No messages to summarize"),
+        }
+
+    formatted = []
+    for msg in messages:
+        role = msg.get("role", "unknown")
+        text = _extract_message_text(msg)
+        formatted.append(f"{role}: {text}")
+
+    conversation = "\n".join(formatted)
+    prompt = SUMMARY_PROMPT_TEMPLATE.format(conversation=conversation)
+
+    try:
+        completion_args = {
+            "model": model,
+            "messages": [{"role": "user", "content": prompt}],
+            "timeout": timeout,
+        }
+
+        response = litellm.completion(**completion_args)
+        summary = response.choices[0].message.content or ""
+        if not summary.strip():
+            return messages[0]
+        summary_msg = "<context_summary message_count='{count}'>{text}</context_summary>"
+        return {
+            "role": "assistant",
+            "content": summary_msg.format(count=len(messages), text=summary),
+        }
+    except Exception:
+        logger.exception("Failed to summarize messages")
+        return messages[0]
+
+
+def _handle_images(messages: list[dict[str, Any]], max_images: int) -> None:
+    image_count = 0
+    for msg in reversed(messages):
+        content = msg.get("content", [])
+        if isinstance(content, list):
+            for item in content:
+                if isinstance(item, dict) and item.get("type") == "image_url":
+                    if image_count >= max_images:
+                        item.update(
+                            {
+                                "type": "text",
+                                "text": "[Previously attached image removed to preserve context]",
+                            }
+                        )
+                    else:
+                        image_count += 1
+
+
+class MemoryCompressor:
+    def __init__(
+        self,
+        max_images: int = 3,
+        model_name: str | None = None,
+        timeout: int = 600,
+    ):
+        self.max_images = max_images
+        self.model_name = model_name or os.getenv("AIPT_LLM", "openai/gpt-5")
+        self.timeout = timeout
+
+        if not self.model_name:
+            raise ValueError("AIPT_LLM environment variable must be set and not empty")
+
+    def compress_history(
+        self,
+        messages: list[dict[str, Any]],
+    ) -> list[dict[str, Any]]:
+        """Compress conversation history to stay within token limits.
+
+        Strategy:
+        1. Handle image limits first
+        2. Keep all system messages
+        3. Keep minimum recent messages
+        4. Summarize older messages when total tokens exceed limit
+
+        The compression preserves:
+        - All system messages unchanged
+        - Most recent messages intact
+        - Critical security context in summaries
+        - Recent images for visual context
+        - Technical details and findings
+        """
+        if not messages:
+            return messages
+
+        _handle_images(messages, self.max_images)
+
+        system_msgs = []
+        regular_msgs = []
+        for msg in messages:
+            if msg.get("role") == "system":
+                system_msgs.append(msg)
+            else:
+                regular_msgs.append(msg)
+
+        recent_msgs = regular_msgs[-MIN_RECENT_MESSAGES:]
+        old_msgs = regular_msgs[:-MIN_RECENT_MESSAGES]
+
+        # Type assertion since we ensure model_name is not None in __init__
+        model_name: str = self.model_name  # type: ignore[assignment]
+
+        total_tokens = sum(
+            _get_message_tokens(msg, model_name) for msg in system_msgs + regular_msgs
+        )
+
+        if total_tokens <= MAX_TOTAL_TOKENS * 0.9:
+            return messages
+
+        compressed = []
+        chunk_size = 10
+        for i in range(0, len(old_msgs), chunk_size):
+            chunk = old_msgs[i : i + chunk_size]
+            summary = _summarize_messages(chunk, model_name, self.timeout)
+            if summary:
+                compressed.append(summary)
+
+        return system_msgs + compressed + recent_msgs

aipt_v2/llm/request_queue.py

@@ -0,0 +1,89 @@
+from __future__ import annotations
+
+import asyncio
+import logging
+import os
+import threading
+import time
+from typing import Any
+
+import litellm
+from litellm import ModelResponse, completion
+from tenacity import retry, retry_if_exception, stop_after_attempt, wait_exponential
+
+
+logger = logging.getLogger(__name__)
+
+
+def should_retry_exception(exception: Exception) -> bool:
+    status_code = None
+
+    if hasattr(exception, "status_code"):
+        status_code = exception.status_code
+    elif hasattr(exception, "response") and hasattr(exception.response, "status_code"):
+        status_code = exception.response.status_code
+
+    if status_code is not None:
+        return bool(litellm._should_retry(status_code))
+    return True
+
+
+class LLMRequestQueue:
+    def __init__(self, max_concurrent: int = 1, delay_between_requests: float = 4.0):
+        rate_limit_delay = os.getenv("LLM_RATE_LIMIT_DELAY")
+        if rate_limit_delay:
+            delay_between_requests = float(rate_limit_delay)
+
+        rate_limit_concurrent = os.getenv("LLM_RATE_LIMIT_CONCURRENT")
+        if rate_limit_concurrent:
+            max_concurrent = int(rate_limit_concurrent)
+
+        self.max_concurrent = max_concurrent
+        self.delay_between_requests = delay_between_requests
+        self._semaphore = threading.BoundedSemaphore(max_concurrent)
+        self._last_request_time = 0.0
+        self._lock = threading.Lock()
+
+    async def make_request(self, completion_args: dict[str, Any]) -> ModelResponse:
+        try:
+            while not self._semaphore.acquire(timeout=0.2):
+                await asyncio.sleep(0.1)
+
+            with self._lock:
+                now = time.time()
+                time_since_last = now - self._last_request_time
+                sleep_needed = max(0, self.delay_between_requests - time_since_last)
+                self._last_request_time = now + sleep_needed
+
+            if sleep_needed > 0:
+                await asyncio.sleep(sleep_needed)
+
+            return await self._reliable_request(completion_args)
+        finally:
+            self._semaphore.release()
+
+    @retry(  # type: ignore[misc]
+        stop=stop_after_attempt(3),
+        wait=wait_exponential(multiplier=8, min=8, max=64),
+        retry=retry_if_exception(should_retry_exception),
+        reraise=True,
+    )
+    async def _reliable_request(self, completion_args: dict[str, Any]) -> ModelResponse:
+        response = completion(**completion_args, stream=False)
+        if isinstance(response, ModelResponse):
+            return response
+        self._raise_unexpected_response()
+        raise RuntimeError("Unreachable code")
+
+    def _raise_unexpected_response(self) -> None:
+        raise RuntimeError("Unexpected response type")
+
+
+_global_queue: LLMRequestQueue | None = None
+
+
+def get_global_queue() -> LLMRequestQueue:
+    global _global_queue  # noqa: PLW0603
+    if _global_queue is None:
+        _global_queue = LLMRequestQueue()
+    return _global_queue

aipt_v2/llm/utils.py

@@ -0,0 +1,89 @@
+from __future__ import annotations
+
+import html
+import re
+from typing import Any
+
+
+def _truncate_to_first_function(content: str) -> str:
+    if not content:
+        return content
+
+    function_starts = [match.start() for match in re.finditer(r"<function=", content)]
+
+    if len(function_starts) >= 2:
+        second_function_start = function_starts[1]
+
+        return content[:second_function_start].rstrip()
+
+    return content
+
+
+def parse_tool_invocations(content: str) -> list[dict[str, Any]] | None:
+    content = _fix_stopword(content)
+
+    tool_invocations: list[dict[str, Any]] = []
+
+    fn_regex_pattern = r"<function=([^>]+)>\n?(.*?)</function>"
+    fn_param_regex_pattern = r"<parameter=([^>]+)>(.*?)</parameter>"
+
+    fn_matches = re.finditer(fn_regex_pattern, content, re.DOTALL)
+
+    for fn_match in fn_matches:
+        fn_name = fn_match.group(1)
+        fn_body = fn_match.group(2)
+
+        param_matches = re.finditer(fn_param_regex_pattern, fn_body, re.DOTALL)
+
+        args = {}
+        for param_match in param_matches:
+            param_name = param_match.group(1)
+            param_value = param_match.group(2).strip()
+
+            param_value = html.unescape(param_value)
+            args[param_name] = param_value
+
+        tool_invocations.append({"toolName": fn_name, "args": args})
+
+    return tool_invocations if tool_invocations else None
+
+
+def _fix_stopword(content: str) -> str:
+    if "<function=" in content and content.count("<function=") == 1:
+        if content.endswith("</"):
+            content = content.rstrip() + "function>"
+        elif not content.rstrip().endswith("</function>"):
+            content = content + "\n</function>"
+    return content
+
+
+def format_tool_call(tool_name: str, args: dict[str, Any]) -> str:
+    xml_parts = [f"<function={tool_name}>"]
+
+    for key, value in args.items():
+        xml_parts.append(f"<parameter={key}>{value}</parameter>")
+
+    xml_parts.append("</function>")
+
+    return "\n".join(xml_parts)
+
+
+def clean_content(content: str) -> str:
+    if not content:
+        return ""
+
+    content = _fix_stopword(content)
+
+    tool_pattern = r"<function=[^>]+>.*?</function>"
+    cleaned = re.sub(tool_pattern, "", content, flags=re.DOTALL)
+
+    hidden_xml_patterns = [
+        r"<inter_agent_message>.*?</inter_agent_message>",
+        r"<agent_completion_report>.*?</agent_completion_report>",
+    ]
+    for pattern in hidden_xml_patterns:
+        cleaned = re.sub(pattern, "", cleaned, flags=re.DOTALL | re.IGNORECASE)
+
+    cleaned = re.sub(r"\n\s*\n", "\n\n", cleaned)
+
+    return cleaned.strip()

aipt_v2/models/__init__.py

@@ -0,0 +1,15 @@
+"""AIPT Data Models"""
+
+from .findings import Finding, Severity, VulnerabilityType
+from .scan_config import ScanConfig, ScanMode
+from .phase_result import PhaseResult, Phase
+
+__all__ = [
+    "Finding",
+    "Severity",
+    "VulnerabilityType",
+    "ScanConfig",
+    "ScanMode",
+    "PhaseResult",
+    "Phase",
+]

aipt_v2/models/findings.py

@@ -0,0 +1,295 @@
+"""
+AIPT Finding Model - Unified vulnerability representation
+
+This model represents vulnerabilities discovered by ANY tool in the pipeline:
+- Traditional scanners (Acunetix, Burp, Nuclei, ZAP)
+- AI-autonomous agents (Strix)
+- Manual exploitation attempts
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from datetime import datetime
+from enum import Enum
+from typing import Any
+import hashlib
+import json
+
+
+class Severity(Enum):
+    """CVSS-aligned severity levels"""
+    CRITICAL = "critical"  # CVSS 9.0-10.0
+    HIGH = "high"          # CVSS 7.0-8.9
+    MEDIUM = "medium"      # CVSS 4.0-6.9
+    LOW = "low"            # CVSS 0.1-3.9
+    INFO = "info"          # CVSS 0.0 / Informational
+
+    @classmethod
+    def from_cvss(cls, score: float) -> "Severity":
+        """Convert CVSS score to severity level"""
+        if score >= 9.0:
+            return cls.CRITICAL
+        elif score >= 7.0:
+            return cls.HIGH
+        elif score >= 4.0:
+            return cls.MEDIUM
+        elif score > 0:
+            return cls.LOW
+        return cls.INFO
+
+    def __lt__(self, other: "Severity") -> bool:
+        order = [self.INFO, self.LOW, self.MEDIUM, self.HIGH, self.CRITICAL]
+        return order.index(self) < order.index(other)
+
+
+class VulnerabilityType(Enum):
+    """OWASP Top 10 aligned vulnerability categories"""
+    # A01:2021 - Broken Access Control
+    IDOR = "idor"
+    BROKEN_ACCESS_CONTROL = "broken_access_control"
+    PRIVILEGE_ESCALATION = "privilege_escalation"
+
+    # A02:2021 - Cryptographic Failures
+    WEAK_CRYPTO = "weak_crypto"
+    SENSITIVE_DATA_EXPOSURE = "sensitive_data_exposure"
+
+    # A03:2021 - Injection
+    SQL_INJECTION = "sql_injection"
+    COMMAND_INJECTION = "command_injection"
+    LDAP_INJECTION = "ldap_injection"
+    XPATH_INJECTION = "xpath_injection"
+    NOSQL_INJECTION = "nosql_injection"
+
+    # A04:2021 - Insecure Design
+    BUSINESS_LOGIC_FLAW = "business_logic_flaw"
+
+    # A05:2021 - Security Misconfiguration
+    MISCONFIGURATION = "misconfiguration"
+    DEFAULT_CREDENTIALS = "default_credentials"
+    DIRECTORY_LISTING = "directory_listing"
+
+    # A06:2021 - Vulnerable Components
+    OUTDATED_COMPONENT = "outdated_component"
+    KNOWN_CVE = "known_cve"
+
+    # A07:2021 - Authentication Failures
+    AUTH_BYPASS = "auth_bypass"
+    WEAK_PASSWORD = "weak_password"
+    SESSION_FIXATION = "session_fixation"
+
+    # A08:2021 - Software Integrity Failures
+    INSECURE_DESERIALIZATION = "insecure_deserialization"
+
+    # A09:2021 - Logging & Monitoring Failures
+    INSUFFICIENT_LOGGING = "insufficient_logging"
+
+    # A10:2021 - SSRF
+    SSRF = "ssrf"
+
+    # Cross-Site Scripting (separate category)
+    XSS_REFLECTED = "xss_reflected"
+    XSS_STORED = "xss_stored"
+    XSS_DOM = "xss_dom"
+
+    # Other
+    OPEN_REDIRECT = "open_redirect"
+    FILE_INCLUSION = "file_inclusion"
+    FILE_UPLOAD = "file_upload"
+    XXE = "xxe"
+    CORS_MISCONFIGURATION = "cors_misconfiguration"
+    CSRF = "csrf"
+    INFORMATION_DISCLOSURE = "information_disclosure"
+    RCE = "rce"
+
+    # Catch-all
+    OTHER = "other"
+
+
+@dataclass
+class Finding:
+    """
+    Unified vulnerability finding from any source
+
+    This is the core data structure that normalizes findings from:
+    - Acunetix (JSON API responses)
+    - Burp Suite (XML/JSON exports)
+    - Nuclei (JSON output)
+    - ZAP (JSON API responses)
+    - Strix (AI agent reports)
+    """
+
+    # Core identification
+    title: str
+    severity: Severity
+    vuln_type: VulnerabilityType
+
+    # Location
+    url: str
+    parameter: str | None = None
+    method: str = "GET"
+
+    # Evidence
+    description: str = ""
+    evidence: str = ""
+    request: str | None = None
+    response: str | None = None
+
+    # Source tracking
+    source: str = "unknown"  # acunetix, burp, nuclei, zap, aipt, manual
+    source_id: str | None = None  # Original ID from source scanner
+
+    # Validation
+    confirmed: bool = False
+    exploited: bool = False
+    poc_command: str | None = None
+
+    # Metadata
+    cvss_score: float | None = None
+    cwe_id: str | None = None
+    cve_ids: list[str] = field(default_factory=list)
+    references: list[str] = field(default_factory=list)
+
+    # Remediation
+    remediation: str = ""
+
+    # Timestamps
+    discovered_at: datetime = field(default_factory=datetime.utcnow)
+
+    # AI-specific fields (for Strix findings)
+    ai_reasoning: str | None = None
+    ai_confidence: float | None = None  # 0.0 to 1.0
+
+    def __post_init__(self):
+        """Generate unique fingerprint for deduplication"""
+        self._fingerprint = self._generate_fingerprint()
+
+    def _generate_fingerprint(self) -> str:
+        """
+        Generate a unique fingerprint for finding deduplication.
+
+        Two findings are considered duplicates if they have the same:
+        - URL (normalized)
+        - Parameter
+        - Vulnerability type
+        """
+        normalized_url = self.url.rstrip("/").lower()
+        data = f"{normalized_url}:{self.parameter}:{self.vuln_type.value}"
+        return hashlib.sha256(data.encode()).hexdigest()[:16]
+
+    @property
+    def fingerprint(self) -> str:
+        return self._fingerprint
+
+    def is_duplicate_of(self, other: "Finding") -> bool:
+        """Check if this finding is a duplicate of another"""
+        return self.fingerprint == other.fingerprint
+
+    def merge_with(self, other: "Finding") -> "Finding":
+        """
+        Merge two duplicate findings, keeping the best evidence from both.
+        Prefers confirmed/exploited findings, higher confidence, more details.
+        """
+        # Prefer the confirmed/exploited finding
+        if other.confirmed and not self.confirmed:
+            base, supplement = other, self
+        elif other.exploited and not self.exploited:
+            base, supplement = other, self
+        else:
+            base, supplement = self, other
+
+        # Merge evidence
+        merged_evidence = base.evidence
+        if supplement.evidence and supplement.evidence not in merged_evidence:
+            merged_evidence = f"{merged_evidence}\n\n--- Additional Evidence ---\n{supplement.evidence}"
+
+        # Merge sources
+        sources = set([base.source, supplement.source])
+        merged_source = ", ".join(sorted(sources))
+
+        # Take highest confidence
+        confidence = max(
+            base.ai_confidence or 0,
+            supplement.ai_confidence or 0
+        ) or None
+
+        return Finding(
+            title=base.title,
+            severity=max(base.severity, other.severity),  # Take highest severity
+            vuln_type=base.vuln_type,
+            url=base.url,
+            parameter=base.parameter,
+            method=base.method,
+            description=base.description or supplement.description,
+            evidence=merged_evidence,
+            request=base.request or supplement.request,
+            response=base.response or supplement.response,
+            source=merged_source,
+            confirmed=base.confirmed or supplement.confirmed,
+            exploited=base.exploited or supplement.exploited,
+            poc_command=base.poc_command or supplement.poc_command,
+            cvss_score=base.cvss_score or supplement.cvss_score,
+            cwe_id=base.cwe_id or supplement.cwe_id,
+            cve_ids=list(set(base.cve_ids + supplement.cve_ids)),
+            references=list(set(base.references + supplement.references)),
+            remediation=base.remediation or supplement.remediation,
+            ai_reasoning=base.ai_reasoning or supplement.ai_reasoning,
+            ai_confidence=confidence,
+        )
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dictionary for JSON serialization"""
+        return {
+            "fingerprint": self.fingerprint,
+            "title": self.title,
+            "severity": self.severity.value,
+            "vuln_type": self.vuln_type.value,
+            "url": self.url,
+            "parameter": self.parameter,
+            "method": self.method,
+            "description": self.description,
+            "evidence": self.evidence,
+            "request": self.request,
+            "response": self.response,
+            "source": self.source,
+            "source_id": self.source_id,
+            "confirmed": self.confirmed,
+            "exploited": self.exploited,
+            "poc_command": self.poc_command,
+            "cvss_score": self.cvss_score,
+            "cwe_id": self.cwe_id,
+            "cve_ids": self.cve_ids,
+            "references": self.references,
+            "remediation": self.remediation,
+            "discovered_at": self.discovered_at.isoformat(),
+            "ai_reasoning": self.ai_reasoning,
+            "ai_confidence": self.ai_confidence,
+        }
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> "Finding":
+        """Create Finding from dictionary"""
+        return cls(
+            title=data["title"],
+            severity=Severity(data["severity"]),
+            vuln_type=VulnerabilityType(data.get("vuln_type", "other")),
+            url=data["url"],
+            parameter=data.get("parameter"),
+            method=data.get("method", "GET"),
+            description=data.get("description", ""),
+            evidence=data.get("evidence", ""),
+            request=data.get("request"),
+            response=data.get("response"),
+            source=data.get("source", "unknown"),
+            source_id=data.get("source_id"),
+            confirmed=data.get("confirmed", False),
+            exploited=data.get("exploited", False),
+            poc_command=data.get("poc_command"),
+            cvss_score=data.get("cvss_score"),
+            cwe_id=data.get("cwe_id"),
+            cve_ids=data.get("cve_ids", []),
+            references=data.get("references", []),
+            remediation=data.get("remediation", ""),
+            discovered_at=datetime.fromisoformat(data["discovered_at"]) if "discovered_at" in data else datetime.utcnow(),
+            ai_reasoning=data.get("ai_reasoning"),
+            ai_confidence=data.get("ai_confidence"),
+        )