pentest_memory.py

"""
Pentest Memory and Learning System.

This module implements a learning system that enables the Red Cell agent to:

1. Remember Previous Findings - Store and recall past vulnerabilities
2. Learn Attack Patterns - Identify what works against specific technologies
3. Adapt Strategies - Improve testing based on historical success
4. Cross-Target Learning - Apply lessons from one target to similar targets
5. Payload Evolution - Track which payloads are most effective

This enables the agent to become more effective over time, learning from
every pentest to improve future assessments.
"""
import hashlib
import json
import os
import re
from collections import defaultdict
from dataclasses import dataclass, field
from datetime import datetime, timedelta
from typing import Any, Dict, List, Optional, Set, Tuple

from openai import AsyncOpenAI
from temporalio import activity

from agentex.lib.utils.logging import make_logger
from agentex.lib import adk
from agentex.types.text_content import TextContent

logger = make_logger(__name__)


async def call_llm(prompt: str, system_prompt: str, temperature: float = 0.3, max_tokens: int = 2000) -> str:
    """Call LLM for analysis."""
    api_key = os.environ.get("OPENAI_API_KEY")
    base_url = os.environ.get("OPENAI_BASE_URL", "https://api.openai.com/v1")
    model = os.environ.get("OPENAI_MODEL", "gpt-4")

    if not api_key:
        raise ValueError("OPENAI_API_KEY not set")

    client = AsyncOpenAI(
        api_key=api_key,
        base_url=base_url,
        timeout=120.0,
    )

    response = await client.chat.completions.create(
        model=model,
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": prompt},
        ],
        temperature=temperature,
        max_tokens=max_tokens,
    )

    return response.choices[0].message.content


# =============================================================================
# MEMORY DATA STRUCTURES
# =============================================================================

@dataclass
class VulnerabilityRecord:
    """Record of a discovered vulnerability."""
    id: str
    target_domain: str
    endpoint: str
    vulnerability_type: str
    payload: str
    severity: str
    confidence: float
    discovered_at: datetime
    technologies: List[str] = field(default_factory=list)
    indicators: List[str] = field(default_factory=list)
    response_sample: str = ""
    verified: bool = False
    false_positive: bool = False
    notes: str = ""

    def to_dict(self) -> Dict[str, Any]:
        return {
            "id": self.id,
            "target_domain": self.target_domain,
            "endpoint": self.endpoint,
            "vulnerability_type": self.vulnerability_type,
            "payload": self.payload,
            "severity": self.severity,
            "confidence": self.confidence,
            "discovered_at": self.discovered_at.isoformat(),
            "technologies": self.technologies,
            "indicators": self.indicators,
            "response_sample": self.response_sample[:500],
            "verified": self.verified,
            "false_positive": self.false_positive,
            "notes": self.notes,
        }

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "VulnerabilityRecord":
        data["discovered_at"] = datetime.fromisoformat(data["discovered_at"])
        return cls(**data)


@dataclass
class PayloadEffectiveness:
    """Tracks effectiveness of a payload."""
    payload: str
    vulnerability_type: str
    success_count: int = 0
    failure_count: int = 0
    technologies: Set[str] = field(default_factory=set)
    last_success: Optional[datetime] = None

    @property
    def success_rate(self) -> float:
        total = self.success_count + self.failure_count
        return self.success_count / total if total > 0 else 0.0

    def record_success(self, technology: str = ""):
        self.success_count += 1
        self.last_success = datetime.utcnow()
        if technology:
            self.technologies.add(technology)

    def record_failure(self):
        self.failure_count += 1


@dataclass
class TechnologyProfile:
    """Profile of vulnerabilities for a technology."""
    technology: str
    vulnerability_counts: Dict[str, int] = field(default_factory=lambda: defaultdict(int))
    effective_payloads: Dict[str, List[str]] = field(default_factory=lambda: defaultdict(list))
    common_endpoints: List[str] = field(default_factory=list)
    last_updated: datetime = field(default_factory=datetime.utcnow)

    def add_vulnerability(self, vuln_type: str, payload: str, endpoint: str):
        self.vulnerability_counts[vuln_type] += 1
        if payload not in self.effective_payloads[vuln_type]:
            self.effective_payloads[vuln_type].append(payload)
        if endpoint not in self.common_endpoints:
            self.common_endpoints.append(endpoint)
        self.last_updated = datetime.utcnow()


class PentestMemory:
    """
    Central memory system for the pentest agent.

    Stores and retrieves:
    - Vulnerability records
    - Payload effectiveness
    - Technology profiles
    - Attack patterns
    """

    def __init__(self, storage_path: Optional[str] = None):
        self.storage_path = storage_path or "/tmp/pentest_memory"
        self.vulnerabilities: Dict[str, VulnerabilityRecord] = {}
        self.payloads: Dict[str, PayloadEffectiveness] = {}
        self.technologies: Dict[str, TechnologyProfile] = {}
        self.attack_patterns: Dict[str, Dict[str, Any]] = {}
        self._load()

    def _get_payload_key(self, payload: str, vuln_type: str) -> str:
        """Generate unique key for payload."""
        return hashlib.md5(f"{vuln_type}:{payload}".encode()).hexdigest()[:12]

    def _load(self):
        """Load memory from storage."""
        try:
            if os.path.exists(f"{self.storage_path}/vulnerabilities.json"):
                with open(f"{self.storage_path}/vulnerabilities.json", "r") as f:
                    data = json.load(f)
                    self.vulnerabilities = {
                        k: VulnerabilityRecord.from_dict(v)
                        for k, v in data.items()
                    }

            if os.path.exists(f"{self.storage_path}/payloads.json"):
                with open(f"{self.storage_path}/payloads.json", "r") as f:
                    data = json.load(f)
                    for k, v in data.items():
                        self.payloads[k] = PayloadEffectiveness(
                            payload=v["payload"],
                            vulnerability_type=v["vulnerability_type"],
                            success_count=v["success_count"],
                            failure_count=v["failure_count"],
                            technologies=set(v.get("technologies", [])),
                            last_success=datetime.fromisoformat(v["last_success"]) if v.get("last_success") else None,
                        )

            if os.path.exists(f"{self.storage_path}/technologies.json"):
                with open(f"{self.storage_path}/technologies.json", "r") as f:
                    data = json.load(f)
                    for k, v in data.items():
                        self.technologies[k] = TechnologyProfile(
                            technology=v["technology"],
                            vulnerability_counts=defaultdict(int, v.get("vulnerability_counts", {})),
                            effective_payloads=defaultdict(list, v.get("effective_payloads", {})),
                            common_endpoints=v.get("common_endpoints", []),
                            last_updated=datetime.fromisoformat(v["last_updated"]) if v.get("last_updated") else datetime.utcnow(),
                        )

            if os.path.exists(f"{self.storage_path}/attack_patterns.json"):
                with open(f"{self.storage_path}/attack_patterns.json", "r") as f:
                    self.attack_patterns = json.load(f)

        except Exception as e:
            logger.warning(f"Failed to load memory: {e}")

    def _save(self):
        """Save memory to storage."""
        try:
            os.makedirs(self.storage_path, exist_ok=True)

            with open(f"{self.storage_path}/vulnerabilities.json", "w") as f:
                json.dump({k: v.to_dict() for k, v in self.vulnerabilities.items()}, f, indent=2)

            with open(f"{self.storage_path}/payloads.json", "w") as f:
                payload_data = {}
                for k, v in self.payloads.items():
                    payload_data[k] = {
                        "payload": v.payload,
                        "vulnerability_type": v.vulnerability_type,
                        "success_count": v.success_count,
                        "failure_count": v.failure_count,
                        "technologies": list(v.technologies),
                        "last_success": v.last_success.isoformat() if v.last_success else None,
                    }
                json.dump(payload_data, f, indent=2)

            with open(f"{self.storage_path}/technologies.json", "w") as f:
                tech_data = {}
                for k, v in self.technologies.items():
                    tech_data[k] = {
                        "technology": v.technology,
                        "vulnerability_counts": dict(v.vulnerability_counts),
                        "effective_payloads": dict(v.effective_payloads),
                        "common_endpoints": v.common_endpoints[:100],
                        "last_updated": v.last_updated.isoformat(),
                    }
                json.dump(tech_data, f, indent=2)

            with open(f"{self.storage_path}/attack_patterns.json", "w") as f:
                json.dump(self.attack_patterns, f, indent=2)

        except Exception as e:
            logger.warning(f"Failed to save memory: {e}")

    def record_vulnerability(
        self,
        target_domain: str,
        endpoint: str,
        vulnerability_type: str,
        payload: str,
        severity: str,
        confidence: float,
        technologies: List[str] = None,
        indicators: List[str] = None,
        response_sample: str = "",
    ) -> VulnerabilityRecord:
        """Record a discovered vulnerability."""
        vuln_id = hashlib.md5(
            f"{target_domain}:{endpoint}:{vulnerability_type}:{payload}".encode()
        ).hexdigest()[:16]

        record = VulnerabilityRecord(
            id=vuln_id,
            target_domain=target_domain,
            endpoint=endpoint,
            vulnerability_type=vulnerability_type,
            payload=payload,
            severity=severity,
            confidence=confidence,
            discovered_at=datetime.utcnow(),
            technologies=technologies or [],
            indicators=indicators or [],
            response_sample=response_sample,
        )

        self.vulnerabilities[vuln_id] = record

        # Update payload effectiveness
        payload_key = self._get_payload_key(payload, vulnerability_type)
        if payload_key not in self.payloads:
            self.payloads[payload_key] = PayloadEffectiveness(
                payload=payload,
                vulnerability_type=vulnerability_type,
            )
        self.payloads[payload_key].record_success(technologies[0] if technologies else "")

        # Update technology profiles
        for tech in (technologies or []):
            if tech not in self.technologies:
                self.technologies[tech] = TechnologyProfile(technology=tech)
            self.technologies[tech].add_vulnerability(vulnerability_type, payload, endpoint)

        self._save()
        return record

    def record_failed_test(self, payload: str, vulnerability_type: str):
        """Record a failed test attempt."""
        payload_key = self._get_payload_key(payload, vulnerability_type)
        if payload_key not in self.payloads:
            self.payloads[payload_key] = PayloadEffectiveness(
                payload=payload,
                vulnerability_type=vulnerability_type,
            )
        self.payloads[payload_key].record_failure()

    def get_effective_payloads(
        self,
        vulnerability_type: str,
        technology: str = None,
        min_success_rate: float = 0.3,
        limit: int = 20,
    ) -> List[Tuple[str, float]]:
        """Get most effective payloads for a vulnerability type."""
        candidates = []

        for payload_eff in self.payloads.values():
            if payload_eff.vulnerability_type != vulnerability_type:
                continue

            if technology and technology not in payload_eff.technologies:
                continue

            if payload_eff.success_rate >= min_success_rate:
                candidates.append((payload_eff.payload, payload_eff.success_rate))

        # Sort by success rate
        candidates.sort(key=lambda x: x[1], reverse=True)
        return candidates[:limit]

    def get_similar_vulnerabilities(
        self,
        target_domain: str = None,
        vulnerability_type: str = None,
        technology: str = None,
        limit: int = 10,
    ) -> List[VulnerabilityRecord]:
        """Get similar vulnerabilities from memory."""
        candidates = []

        for vuln in self.vulnerabilities.values():
            score = 0

            if target_domain and vuln.target_domain == target_domain:
                score += 3

            if vulnerability_type and vuln.vulnerability_type == vulnerability_type:
                score += 2

            if technology and technology in vuln.technologies:
                score += 1

            if score > 0:
                candidates.append((vuln, score))

        candidates.sort(key=lambda x: x[1], reverse=True)
        return [c[0] for c in candidates[:limit]]

    def get_technology_insights(self, technology: str) -> Dict[str, Any]:
        """Get insights about a technology's vulnerabilities."""
        if technology not in self.technologies:
            return {"technology": technology, "insights": "No data available"}

        profile = self.technologies[technology]

        # Get top vulnerability types
        sorted_vulns = sorted(
            profile.vulnerability_counts.items(),
            key=lambda x: x[1],
            reverse=True,
        )

        return {
            "technology": technology,
            "total_vulnerabilities": sum(profile.vulnerability_counts.values()),
            "top_vulnerability_types": sorted_vulns[:5],
            "effective_payloads": {
                k: v[:5] for k, v in profile.effective_payloads.items()
            },
            "common_endpoints": profile.common_endpoints[:10],
            "last_updated": profile.last_updated.isoformat(),
        }

    def record_attack_pattern(
        self,
        pattern_name: str,
        description: str,
        steps: List[str],
        success_rate: float,
        applicable_technologies: List[str],
    ):
        """Record a successful attack pattern."""
        self.attack_patterns[pattern_name] = {
            "name": pattern_name,
            "description": description,
            "steps": steps,
            "success_rate": success_rate,
            "applicable_technologies": applicable_technologies,
            "recorded_at": datetime.utcnow().isoformat(),
        }
        self._save()

    def get_applicable_patterns(
        self,
        technologies: List[str],
        limit: int = 5,
    ) -> List[Dict[str, Any]]:
        """Get attack patterns applicable to given technologies."""
        candidates = []

        for pattern in self.attack_patterns.values():
            applicable = pattern.get("applicable_technologies", [])
            overlap = len(set(applicable) & set(technologies))

            if overlap > 0:
                candidates.append((pattern, overlap, pattern.get("success_rate", 0)))

        # Sort by overlap and success rate
        candidates.sort(key=lambda x: (x[1], x[2]), reverse=True)
        return [c[0] for c in candidates[:limit]]

    def get_statistics(self) -> Dict[str, Any]:
        """Get memory statistics."""
        return {
            "total_vulnerabilities": len(self.vulnerabilities),
            "total_payloads_tracked": len(self.payloads),
            "technologies_profiled": len(self.technologies),
            "attack_patterns_recorded": len(self.attack_patterns),
            "vulnerabilities_by_type": defaultdict(int, {
                v.vulnerability_type: 1 for v in self.vulnerabilities.values()
            }),
            "top_effective_payloads": [
                (p.payload[:50], p.success_rate)
                for p in sorted(self.payloads.values(), key=lambda x: x.success_rate, reverse=True)[:10]
            ],
        }


# Global memory instance
_pentest_memory = PentestMemory()


# =============================================================================
# MEMORY ACTIVITIES
# =============================================================================

@activity.defn(name="store_vulnerability_finding_activity")
async def store_vulnerability_finding_activity(
    target_domain: str,
    endpoint: str,
    vulnerability_type: str,
    payload: str,
    severity: str,
    confidence: float,
    technologies: List[str],
    indicators: List[str],
    response_sample: str,
    task_id: str,
    trace_id: str,
) -> Dict[str, Any]:
    """
    Store a vulnerability finding in memory.

    This activity:
    1. Records the vulnerability
    2. Updates payload effectiveness
    3. Updates technology profiles
    4. Enables future learning
    """
    logger.info(f"Storing vulnerability: {vulnerability_type} on {endpoint}")

    activity.heartbeat("Storing vulnerability in memory")

    record = _pentest_memory.record_vulnerability(
        target_domain=target_domain,
        endpoint=endpoint,
        vulnerability_type=vulnerability_type,
        payload=payload,
        severity=severity,
        confidence=confidence,
        technologies=technologies,
        indicators=indicators,
        response_sample=response_sample,
    )

    return {
        "vulnerability_id": record.id,
        "stored": True,
        "memory_stats": _pentest_memory.get_statistics(),
    }


@activity.defn(name="get_learned_payloads_activity")
async def get_learned_payloads_activity(
    vulnerability_type: str,
    technologies: List[str],
    task_id: str,
    trace_id: str,
) -> Dict[str, Any]:
    """
    Get effective payloads learned from previous pentests.

    This activity:
    1. Queries memory for effective payloads
    2. Filters by technology if specified
    3. Returns ranked payloads by success rate
    """
    logger.info(f"Getting learned payloads for {vulnerability_type}")

    activity.heartbeat("Querying payload memory")

    all_payloads = []

    # Get payloads for each technology
    for tech in technologies:
        payloads = _pentest_memory.get_effective_payloads(
            vulnerability_type=vulnerability_type,
            technology=tech,
            min_success_rate=0.2,
            limit=10,
        )
        all_payloads.extend(payloads)

    # Also get general payloads
    general_payloads = _pentest_memory.get_effective_payloads(
        vulnerability_type=vulnerability_type,
        min_success_rate=0.3,
        limit=20,
    )
    all_payloads.extend(general_payloads)

    # Deduplicate and sort
    seen = set()
    unique_payloads = []
    for payload, rate in sorted(all_payloads, key=lambda x: x[1], reverse=True):
        if payload not in seen:
            seen.add(payload)
            unique_payloads.append({"payload": payload, "success_rate": rate})

    # Notify about learned payloads
    if task_id and unique_payloads:
        await adk.messages.create(
            task_id=task_id,
            content=TextContent(
                author="agent",
                content=f"""### 🧠 Learned Payloads Retrieved

**Vulnerability Type:** {vulnerability_type}
**Technologies:** {', '.join(technologies)}
**Effective Payloads Found:** {len(unique_payloads)}

**Top Payloads by Success Rate:**
{chr(10).join([f"- `{p['payload'][:50]}...` ({p['success_rate']:.0%})" for p in unique_payloads[:5]])}

Using learned payloads from previous pentests for improved effectiveness.""",
            ),
            trace_id=trace_id,
        )

    return {
        "vulnerability_type": vulnerability_type,
        "technologies": technologies,
        "payloads": unique_payloads[:30],
        "total_found": len(unique_payloads),
    }


@activity.defn(name="get_technology_insights_activity")
async def get_technology_insights_activity(
    technologies: List[str],
    task_id: str,
    trace_id: str,
) -> Dict[str, Any]:
    """
    Get insights about technologies from memory.

    This activity:
    1. Retrieves vulnerability patterns for technologies
    2. Gets effective payloads
    3. Identifies common vulnerable endpoints
    """
    logger.info(f"Getting insights for technologies: {technologies}")

    activity.heartbeat("Querying technology insights")

    insights = {}
    for tech in technologies:
        insights[tech] = _pentest_memory.get_technology_insights(tech)

    # Get applicable attack patterns
    patterns = _pentest_memory.get_applicable_patterns(technologies, limit=5)

    # Notify about insights
    if task_id:
        tech_summary = "\n".join([
            f"- **{tech}**: {ins.get('total_vulnerabilities', 0)} vulns, Top: {', '.join([v[0] for v in ins.get('top_vulnerability_types', [])[:3]])}"
            for tech, ins in insights.items()
            if ins.get('total_vulnerabilities', 0) > 0
        ])

        await adk.messages.create(
            task_id=task_id,
            content=TextContent(
                author="agent",
                content=f"""### 📊 Technology Insights from Memory

**Technologies Analyzed:** {len(technologies)}

**Vulnerability History:**
{tech_summary if tech_summary else "No historical data for these technologies"}

**Applicable Attack Patterns:** {len(patterns)}
{chr(10).join([f"- {p['name']}: {p['description'][:60]}..." for p in patterns[:3]])}

Using historical knowledge to guide testing strategy.""",
            ),
            trace_id=trace_id,
        )

    return {
        "technologies": technologies,
        "insights": insights,
        "attack_patterns": patterns,
    }


@activity.defn(name="get_similar_findings_activity")
async def get_similar_findings_activity(
    target_domain: str,
    vulnerability_type: str,
    technologies: List[str],
    task_id: str,
    trace_id: str,
) -> Dict[str, Any]:
    """
    Get similar vulnerability findings from memory.

    This activity:
    1. Searches for similar vulnerabilities
    2. Returns relevant findings for context
    3. Helps identify patterns
    """
    logger.info(f"Getting similar findings for {target_domain}")

    activity.heartbeat("Searching memory for similar findings")

    similar = []

    # Search by domain
    domain_similar = _pentest_memory.get_similar_vulnerabilities(
        target_domain=target_domain,
        limit=5,
    )
    similar.extend(domain_similar)

    # Search by vulnerability type
    type_similar = _pentest_memory.get_similar_vulnerabilities(
        vulnerability_type=vulnerability_type,
        limit=5,
    )
    similar.extend(type_similar)

    # Search by technology
    for tech in technologies[:3]:
        tech_similar = _pentest_memory.get_similar_vulnerabilities(
            technology=tech,
            limit=3,
        )
        similar.extend(tech_similar)

    # Deduplicate
    seen_ids = set()
    unique_similar = []
    for vuln in similar:
        if vuln.id not in seen_ids:
            seen_ids.add(vuln.id)
            unique_similar.append(vuln.to_dict())

    return {
        "target_domain": target_domain,
        "similar_findings": unique_similar[:15],
        "total_found": len(unique_similar),
    }


@activity.defn(name="analyze_learning_opportunities_activity")
async def analyze_learning_opportunities_activity(
    recent_findings: List[Dict[str, Any]],
    task_id: str,
    trace_id: str,
) -> Dict[str, Any]:
    """
    Analyze recent findings to identify learning opportunities.

    This activity:
    1. Analyzes patterns in recent findings
    2. Identifies new attack patterns
    3. Suggests improvements to testing strategy
    """
    logger.info(f"Analyzing {len(recent_findings)} findings for learning")

    activity.heartbeat("Analyzing learning opportunities")

    if not recent_findings:
        return {"learning_opportunities": [], "patterns_identified": []}

    # Use AI to analyze patterns
    analysis_prompt = f"""Analyze these recent vulnerability findings to identify patterns and learning opportunities:

{json.dumps(recent_findings[:20], indent=2)}

Identify:
1. Common patterns across findings
2. New attack techniques that were effective
3. Technologies that are particularly vulnerable
4. Recommendations for future testing

Return JSON:
{{
  "patterns": [
    {{
      "name": "Pattern name",
      "description": "What the pattern is",
      "applicable_to": ["technologies"],
      "success_indicators": ["what indicates this works"]
    }}
  ],
  "new_techniques": [
    {{
      "name": "Technique name",
      "description": "How it works",
      "payload_template": "Example payload"
    }}
  ],
  "vulnerable_technologies": [
    {{
      "technology": "name",
      "common_vulnerabilities": ["types"],
      "recommended_focus": "what to test"
    }}
  ],
  "recommendations": ["List of recommendations for future testing"]
}}"""

    try:
        response = await call_llm(
            analysis_prompt,
            "You are a security expert analyzing penetration testing results to identify patterns and improve future testing.",
            temperature=0.3,
        )

        json_match = re.search(r'\{[\s\S]*\}', response)
        if json_match:
            analysis = json.loads(json_match.group())
        else:
            analysis = {}

        # Store identified patterns
        for pattern in analysis.get("patterns", []):
            _pentest_memory.record_attack_pattern(
                pattern_name=pattern.get("name", "Unknown"),
                description=pattern.get("description", ""),
                steps=pattern.get("success_indicators", []),
                success_rate=0.7,  # Default for newly identified patterns
                applicable_technologies=pattern.get("applicable_to", []),
            )

        # Notify about learning
        if task_id:
            await adk.messages.create(
                task_id=task_id,
                content=TextContent(
                    author="agent",
                    content=f"""### 📚 Learning Analysis Complete

**Findings Analyzed:** {len(recent_findings)}

**Patterns Identified:** {len(analysis.get('patterns', []))}
{chr(10).join([f"- {p['name']}: {p['description'][:60]}..." for p in analysis.get('patterns', [])[:3]])}

**New Techniques Discovered:** {len(analysis.get('new_techniques', []))}

**Vulnerable Technologies:**
{chr(10).join([f"- {t['technology']}: {', '.join(t.get('common_vulnerabilities', [])[:3])}" for t in analysis.get('vulnerable_technologies', [])[:3]])}

**Recommendations:**
{chr(10).join([f"- {r}" for r in analysis.get('recommendations', [])[:5]])}

Knowledge has been stored for future pentests.""",
                ),
                trace_id=trace_id,
            )

        return {
            "findings_analyzed": len(recent_findings),
            "analysis": analysis,
            "patterns_stored": len(analysis.get("patterns", [])),
        }

    except Exception as e:
        logger.error(f"Learning analysis failed: {e}")
        return {
            "error": str(e),
            "learning_opportunities": [],
        }


@activity.defn(name="get_memory_statistics_activity")
async def get_memory_statistics_activity(
    task_id: str,
    trace_id: str,
) -> Dict[str, Any]:
    """
    Get statistics about the pentest memory.

    This activity:
    1. Returns memory statistics
    2. Shows learning progress
    3. Identifies knowledge gaps
    """
    logger.info("Getting memory statistics")

    activity.heartbeat("Gathering memory statistics")

    stats = _pentest_memory.get_statistics()

    # Notify about statistics
    if task_id:
        await adk.messages.create(
            task_id=task_id,
            content=TextContent(
                author="agent",
                content=f"""### 🧠 Pentest Memory Statistics

**Knowledge Base:**
- Vulnerabilities Recorded: {stats['total_vulnerabilities']}
- Payloads Tracked: {stats['total_payloads_tracked']}
- Technologies Profiled: {stats['technologies_profiled']}
- Attack Patterns: {stats['attack_patterns_recorded']}

**Top Effective Payloads:**
{chr(10).join([f"- `{p[0]}...` ({p[1]:.0%})" for p in stats.get('top_effective_payloads', [])[:5]])}

The agent is continuously learning and improving from each pentest.""",
            ),
            trace_id=trace_id,
        )

    return stats


@activity.defn(name="apply_learned_strategy_activity")
async def apply_learned_strategy_activity(
    target_domain: str,
    technologies: List[str],
    endpoints: List[str],
    task_id: str,
    trace_id: str,
) -> Dict[str, Any]:
    """
    Apply learned strategies to a new target.

    This activity:
    1. Retrieves relevant knowledge from memory
    2. Generates a customized testing strategy
    3. Prioritizes tests based on historical success
    """
    logger.info(f"Applying learned strategy for {target_domain}")

    activity.heartbeat("Generating learned strategy")

    # Get technology insights
    tech_insights = {}
    for tech in technologies:
        tech_insights[tech] = _pentest_memory.get_technology_insights(tech)

    # Get applicable attack patterns
    patterns = _pentest_memory.get_applicable_patterns(technologies, limit=10)

    # Get effective payloads for common vulnerability types
    vuln_types = ["sqli", "xss", "ssrf", "path_traversal", "cmd_injection"]
    effective_payloads = {}

    for vuln_type in vuln_types:
        payloads = _pentest_memory.get_effective_payloads(
            vulnerability_type=vuln_type,
            min_success_rate=0.3,
            limit=10,
        )
        if payloads:
            effective_payloads[vuln_type] = [p[0] for p in payloads]

    # Generate strategy using AI
    strategy_prompt = f"""Based on historical pentest data, generate a testing strategy for:

Target: {target_domain}
Technologies: {', '.join(technologies)}
Endpoints: {len(endpoints)}

Technology Insights:
{json.dumps(tech_insights, indent=2, default=str)}

Known Attack Patterns:
{json.dumps(patterns[:5], indent=2)}

Effective Payloads:
{json.dumps(effective_payloads, indent=2)}

Generate a prioritized testing strategy that:
1. Focuses on historically vulnerable areas
2. Uses proven effective payloads
3. Applies relevant attack patterns
4. Prioritizes high-value endpoints

Return JSON:
{{
  "priority_tests": [
    {{
      "test_name": "...",
      "target_endpoints": ["patterns to match"],
      "vulnerability_types": ["..."],
      "payloads_to_use": ["..."],
      "reasoning": "why this is prioritized"
    }}
  ],
  "attack_sequence": ["ordered list of attack phases"],
  "focus_areas": ["specific areas to focus on"],
  "estimated_effectiveness": 0.0-1.0
}}"""

    try:
        response = await call_llm(
            strategy_prompt,
            "You are a security expert creating a testing strategy based on historical pentest data.",
            temperature=0.4,
        )

        json_match = re.search(r'\{[\s\S]*\}', response)
        if json_match:
            strategy = json.loads(json_match.group())
        else:
            strategy = {"priority_tests": [], "attack_sequence": [], "focus_areas": []}

        # Notify about strategy
        if task_id:
            await adk.messages.create(
                task_id=task_id,
                content=TextContent(
                    author="agent",
                    content=f"""### 🎯 Learned Strategy Applied

**Target:** `{target_domain}`
**Technologies:** {', '.join(technologies)}

**Priority Tests:**
{chr(10).join([f"- {t['test_name']}: {t['reasoning'][:60]}..." for t in strategy.get('priority_tests', [])[:5]])}

**Attack Sequence:**
{chr(10).join([f"{i+1}. {phase}" for i, phase in enumerate(strategy.get('attack_sequence', [])[:5])])}

**Focus Areas:**
{chr(10).join([f"- {area}" for area in strategy.get('focus_areas', [])[:5]])}

**Estimated Effectiveness:** {strategy.get('estimated_effectiveness', 0):.0%}

Strategy generated from {_pentest_memory.get_statistics()['total_vulnerabilities']} historical findings.""",
                ),
                trace_id=trace_id,
            )

        return {
            "target_domain": target_domain,
            "strategy": strategy,
            "effective_payloads": effective_payloads,
            "applicable_patterns": patterns,
        }

    except Exception as e:
        logger.error(f"Strategy generation failed: {e}")
        return {
            "target_domain": target_domain,
            "error": str(e),
            "strategy": {},
        }