continuous_pentest_workflow.py

"""
Continuous Pentesting Workflow.

This workflow orchestrates continuous security assessment that:
1. Never sleeps - runs continuously monitoring for changes
2. Discovers zero-day vulnerabilities through creative AI reasoning
3. Tests faster and more thoroughly than human pentesters
4. Learns from every assessment to improve over time
5. Dramatically reduces costs through automation

This is the "mind-blowing" implementation that combines all advanced
capabilities into a cohesive continuous pentesting system.
"""
import asyncio
from datetime import timedelta
from typing import Any, Dict, List, Optional

from temporalio import workflow
from temporalio.common import RetryPolicy

from agentex.lib.utils.logging import make_logger

logger = make_logger(__name__)

# Activity retry policy
DEFAULT_RETRY_POLICY = RetryPolicy(
    initial_interval=timedelta(seconds=5),
    maximum_interval=timedelta(minutes=5),
    maximum_attempts=3,
    backoff_coefficient=2.0,
)


@workflow.defn(name="ContinuousPentestWorkflow")
class ContinuousPentestWorkflow:
    """
    Continuous Pentesting Workflow that never sleeps.

    This workflow implements a continuous security assessment loop that:
    1. Monitors attack surface for changes
    2. Prioritizes targets based on risk and history
    3. Runs parallel vulnerability scans
    4. Discovers zero-day vulnerabilities through AI reasoning
    5. Learns from findings to improve future assessments
    6. Reports findings in real-time
    """

    def __init__(self):
        self.target_domain: str = ""
        self.task_id: str = ""
        self.trace_id: str = ""
        self.is_running: bool = False
        self.scan_interval_minutes: int = 60
        self.findings: List[Dict[str, Any]] = []
        self.scan_count: int = 0
        self.last_scan_time: Optional[str] = None
        self.technologies: List[str] = []
        self.endpoints: List[str] = []
        self.attack_surface_state: Dict[str, Any] = {}

    @workflow.signal
    async def stop_continuous_scan(self):
        """Signal to stop the continuous scanning loop."""
        self.is_running = False
        workflow.logger.info("Received stop signal - will stop after current scan")

    @workflow.signal
    async def update_scan_interval(self, interval_minutes: int):
        """Update the scan interval."""
        self.scan_interval_minutes = interval_minutes
        workflow.logger.info(f"Updated scan interval to {interval_minutes} minutes")

    @workflow.signal
    async def add_target(self, target: str):
        """Add a new target to the scan scope."""
        if target not in self.endpoints:
            self.endpoints.append(target)
            workflow.logger.info(f"Added new target: {target}")

    @workflow.query
    def get_status(self) -> Dict[str, Any]:
        """Get current workflow status."""
        return {
            "target_domain": self.target_domain,
            "is_running": self.is_running,
            "scan_count": self.scan_count,
            "last_scan_time": self.last_scan_time,
            "total_findings": len(self.findings),
            "endpoints_monitored": len(self.endpoints),
            "technologies": self.technologies,
            "scan_interval_minutes": self.scan_interval_minutes,
        }

    @workflow.query
    def get_findings(self) -> List[Dict[str, Any]]:
        """Get all findings."""
        return self.findings

    @workflow.run
    async def run(
        self,
        target_domain: str,
        task_id: str,
        trace_id: str,
        scan_interval_minutes: int = 60,
        enable_zero_day_discovery: bool = True,
        enable_learning: bool = True,
        max_concurrent_tests: int = 20,
        requests_per_second: float = 50.0,
    ) -> Dict[str, Any]:
        """
        Run the continuous pentesting workflow.

        Args:
            target_domain: The target domain to assess
            task_id: Task ID for messaging
            trace_id: Trace ID for logging
            scan_interval_minutes: Minutes between full scans
            enable_zero_day_discovery: Enable AI-driven zero-day discovery
            enable_learning: Enable learning from findings
            max_concurrent_tests: Maximum concurrent test tasks
            requests_per_second: Rate limit for requests
        """
        self.target_domain = target_domain
        self.task_id = task_id
        self.trace_id = trace_id
        self.scan_interval_minutes = scan_interval_minutes
        self.is_running = True

        workflow.logger.info(f"Starting continuous pentest for {target_domain}")

        # Phase 1: Initial Discovery and Learning
        await self._run_initial_discovery()

        # Phase 2: Apply Learned Strategy
        if enable_learning:
            await self._apply_learned_strategy()

        # Phase 3: Continuous Scanning Loop
        while self.is_running:
            self.scan_count += 1
            workflow.logger.info(f"Starting scan cycle {self.scan_count}")

            # Check for attack surface changes
            changes = await self._detect_attack_surface_changes()

            if changes.get("has_changes", False) or self.scan_count == 1:
                # Run comprehensive scan on changes
                await self._run_comprehensive_scan(
                    max_concurrent=max_concurrent_tests,
                    requests_per_second=requests_per_second,
                )

                # Zero-day discovery
                if enable_zero_day_discovery:
                    await self._run_zero_day_discovery()

                # Learn from findings
                if enable_learning and self.findings:
                    await self._learn_from_findings()

            # Update last scan time
            self.last_scan_time = workflow.now().isoformat()

            # Wait for next scan interval
            await workflow.sleep(timedelta(minutes=self.scan_interval_minutes))

        # Final report
        return await self._generate_final_report()

    async def _run_initial_discovery(self):
        """Run initial asset and endpoint discovery."""
        workflow.logger.info("Running initial discovery")

        # Discover assets
        discovery_result = await workflow.execute_activity(
            "continuous_asset_discovery_activity",
            args=[
                self.target_domain,
                {},  # Empty previous state
                self.task_id,
                self.trace_id,
            ],
            start_to_close_timeout=timedelta(minutes=30),
            retry_policy=DEFAULT_RETRY_POLICY,
        )

        self.endpoints = discovery_result.get("endpoints", [])
        self.technologies = discovery_result.get("technologies", [])
        self.attack_surface_state = discovery_result.get("current_state", {})

        workflow.logger.info(f"Discovered {len(self.endpoints)} endpoints")

    async def _apply_learned_strategy(self):
        """Apply learned strategy from previous pentests."""
        workflow.logger.info("Applying learned strategy")

        strategy_result = await workflow.execute_activity(
            "apply_learned_strategy_activity",
            args=[
                self.target_domain,
                self.technologies,
                self.endpoints,
                self.task_id,
                self.trace_id,
            ],
            start_to_close_timeout=timedelta(minutes=10),
            retry_policy=DEFAULT_RETRY_POLICY,
        )

        # Store effective payloads for use in scanning
        self._learned_payloads = strategy_result.get("effective_payloads", {})
        self._attack_patterns = strategy_result.get("applicable_patterns", [])

    async def _detect_attack_surface_changes(self) -> Dict[str, Any]:
        """Detect changes in the attack surface."""
        workflow.logger.info("Detecting attack surface changes")

        change_result = await workflow.execute_activity(
            "endpoint_change_detection_activity",
            args=[
                self.target_domain,
                self.endpoints,
                self.attack_surface_state,
                self.task_id,
                self.trace_id,
            ],
            start_to_close_timeout=timedelta(minutes=15),
            retry_policy=DEFAULT_RETRY_POLICY,
        )

        # Update state
        if change_result.get("has_changes"):
            self.endpoints = change_result.get("current_endpoints", self.endpoints)
            self.attack_surface_state = change_result.get("new_state", self.attack_surface_state)

        return change_result

    async def _run_comprehensive_scan(
        self,
        max_concurrent: int,
        requests_per_second: float,
    ):
        """Run comprehensive parallel vulnerability scan."""
        workflow.logger.info(f"Running comprehensive scan on {len(self.endpoints)} endpoints")

        # Prioritize endpoints
        priority_result = await workflow.execute_activity(
            "prioritize_endpoints_activity",
            args=[
                self.endpoints,
                self.technologies,
                self.findings,
                self.task_id,
                self.trace_id,
            ],
            start_to_close_timeout=timedelta(minutes=5),
            retry_policy=DEFAULT_RETRY_POLICY,
        )

        prioritized_endpoints = priority_result.get("prioritized_endpoints", self.endpoints)

        # Get learned payloads or use defaults
        payloads_per_type = getattr(self, "_learned_payloads", {})
        if not payloads_per_type:
            payloads_per_type = {
                "sqli": ["'", "\"", "1' OR '1'='1", "1; DROP TABLE users--", "' UNION SELECT NULL--"],
                "xss": ["<script>alert(1)</script>", "<img src=x onerror=alert(1)>", "javascript:alert(1)"],
                "path_traversal": ["../../../etc/passwd", "..\\..\\..\\windows\\system32\\config\\sam"],
                "ssrf": ["http://127.0.0.1", "http://localhost:22", "http://169.254.169.254"],
                "cmd_injection": ["; ls", "| cat /etc/passwd", "`id`", "$(whoami)"],
            }

        # Run parallel scan
        scan_result = await workflow.execute_activity(
            "parallel_vulnerability_scan_activity",
            args=[
                prioritized_endpoints[:500],  # Limit endpoints per scan
                list(payloads_per_type.keys()),
                payloads_per_type,
                max_concurrent,
                requests_per_second,
                self.task_id,
                self.trace_id,
            ],
            start_to_close_timeout=timedelta(hours=2),
            retry_policy=DEFAULT_RETRY_POLICY,
            heartbeat_timeout=timedelta(minutes=5),
        )

        # Store findings
        new_findings = scan_result.get("vulnerabilities", [])
        for finding in new_findings:
            finding["scan_cycle"] = self.scan_count
            finding["discovered_at"] = workflow.now().isoformat()
            self.findings.append(finding)

            # Store in memory for learning
            await workflow.execute_activity(
                "store_vulnerability_finding_activity",
                args=[
                    self.target_domain,
                    finding.get("endpoint", ""),
                    finding.get("test_type", ""),
                    finding.get("payload", ""),
                    "high" if finding.get("vulnerable") else "medium",
                    0.8,
                    self.technologies,
                    finding.get("indicators", []),
                    finding.get("response_sample", ""),
                    self.task_id,
                    self.trace_id,
                ],
                start_to_close_timeout=timedelta(minutes=2),
                retry_policy=DEFAULT_RETRY_POLICY,
            )

        workflow.logger.info(f"Found {len(new_findings)} vulnerabilities in scan cycle {self.scan_count}")

    async def _run_zero_day_discovery(self):
        """Run zero-day discovery through creative AI reasoning."""
        workflow.logger.info("Running zero-day discovery")

        # Behavioral anomaly detection
        anomaly_result = await workflow.execute_activity(
            "behavioral_anomaly_detection_activity",
            args=[
                f"https://{self.target_domain}",
                self.endpoints[:100],  # Sample endpoints
                self.task_id,
                self.trace_id,
            ],
            start_to_close_timeout=timedelta(minutes=30),
            retry_policy=DEFAULT_RETRY_POLICY,
            heartbeat_timeout=timedelta(minutes=5),
        )

        # Store anomaly findings
        for anomaly in anomaly_result.get("anomalies", []):
            if anomaly.get("ai_analysis", {}).get("likely_vulnerability"):
                self.findings.append({
                    "type": "zero_day_candidate",
                    "source": "behavioral_anomaly",
                    "scan_cycle": self.scan_count,
                    **anomaly,
                })

        # Semantic vulnerability reasoning
        semantic_result = await workflow.execute_activity(
            "semantic_vulnerability_reasoning_activity",
            args=[
                f"https://{self.target_domain}",
                self.endpoints[:50],
                self.technologies,
                self.findings[-20:],  # Recent findings as context
                self.task_id,
                self.trace_id,
            ],
            start_to_close_timeout=timedelta(minutes=20),
            retry_policy=DEFAULT_RETRY_POLICY,
        )

        # Store semantic findings
        for vuln in semantic_result.get("potential_vulnerabilities", []):
            if vuln.get("confidence", 0) > 0.6:
                self.findings.append({
                    "type": "zero_day_candidate",
                    "source": "semantic_reasoning",
                    "scan_cycle": self.scan_count,
                    **vuln,
                })

        # Generate and test novel attack vectors
        novel_result = await workflow.execute_activity(
            "generate_novel_attack_vectors_activity",
            args=[
                f"https://{self.target_domain}",
                self.technologies,
                self.endpoints[:30],
                ["sqli", "xss", "ssrf", "path_traversal", "cmd_injection"],
                self.task_id,
                self.trace_id,
            ],
            start_to_close_timeout=timedelta(minutes=15),
            retry_policy=DEFAULT_RETRY_POLICY,
        )

        # Execute top novel attacks
        for attack in novel_result.get("novel_attacks", [])[:5]:
            attack_result = await workflow.execute_activity(
                "execute_novel_attack_activity",
                args=[
                    f"https://{self.target_domain}",
                    attack,
                    self.task_id,
                    self.trace_id,
                ],
                start_to_close_timeout=timedelta(minutes=5),
                retry_policy=DEFAULT_RETRY_POLICY,
            )

            if attack_result.get("success"):
                self.findings.append({
                    "type": "zero_day_candidate",
                    "source": "novel_attack",
                    "scan_cycle": self.scan_count,
                    **attack_result,
                })

        workflow.logger.info(f"Zero-day discovery complete for cycle {self.scan_count}")

    async def _learn_from_findings(self):
        """Learn from findings to improve future scans."""
        workflow.logger.info("Learning from findings")

        # Get recent findings for analysis
        recent_findings = [
            f for f in self.findings
            if f.get("scan_cycle") == self.scan_count
        ]

        if recent_findings:
            await workflow.execute_activity(
                "analyze_learning_opportunities_activity",
                args=[
                    recent_findings,
                    self.task_id,
                    self.trace_id,
                ],
                start_to_close_timeout=timedelta(minutes=10),
                retry_policy=DEFAULT_RETRY_POLICY,
            )

    async def _generate_final_report(self) -> Dict[str, Any]:
        """Generate final report of all findings."""
        workflow.logger.info("Generating final report")

        # Get memory statistics
        memory_stats = await workflow.execute_activity(
            "get_memory_statistics_activity",
            args=[
                self.task_id,
                self.trace_id,
            ],
            start_to_close_timeout=timedelta(minutes=2),
            retry_policy=DEFAULT_RETRY_POLICY,
        )

        # Categorize findings
        by_type = {}
        by_severity = {"critical": [], "high": [], "medium": [], "low": []}
        zero_day_candidates = []

        for finding in self.findings:
            # By type
            ftype = finding.get("test_type") or finding.get("type", "unknown")
            if ftype not in by_type:
                by_type[ftype] = []
            by_type[ftype].append(finding)

            # By severity
            severity = finding.get("severity", "medium")
            if severity in by_severity:
                by_severity[severity].append(finding)

            # Zero-day candidates
            if finding.get("type") == "zero_day_candidate":
                zero_day_candidates.append(finding)

        return {
            "target_domain": self.target_domain,
            "scan_cycles_completed": self.scan_count,
            "total_findings": len(self.findings),
            "findings_by_type": {k: len(v) for k, v in by_type.items()},
            "findings_by_severity": {k: len(v) for k, v in by_severity.items()},
            "zero_day_candidates": len(zero_day_candidates),
            "endpoints_monitored": len(self.endpoints),
            "technologies_detected": self.technologies,
            "memory_statistics": memory_stats,
            "all_findings": self.findings,
        }