scanner_engine.py

"""
Core scanning engine that orchestrates crawling, detection, and AI classification.
Provides real-time progress callbacks for the Streamlit UI.
"""

import logging
import time
from typing import List, Dict, Any, Callable, Optional

from crawler import WebCrawler
from detectors.xss_detector import detect_xss
from detectors.sql_detector import detect_sqli
from detectors.header_detector import detect_missing_headers
from detectors.redirect_detector import detect_open_redirect
from detectors.directory_detector import detect_directories
from ai.vulnerability_ai import VulnerabilityClassifier
from utils.request_manager import RequestManager

logger = logging.getLogger(__name__)


class ScannerEngine:
    """
    Orchestrates the full vulnerability scanning workflow:
    1. Crawl the target
    2. Run detectors on each discovered page
    3. Classify severity via AI
    4. Aggregate and deduplicate results
    """

    def __init__(
        self,
        target_url: str,
        max_pages: int = 20,
        scan_xss: bool = True,
        scan_sqli: bool = True,
        scan_headers: bool = True,
        scan_redirect: bool = True,
        scan_directories: bool = True,
        request_timeout: int = 10,
        request_delay: float = 0.3,
    ):
        self.target_url = target_url.rstrip("/")
        self.max_pages = max_pages
        self.scan_xss = scan_xss
        self.scan_sqli = scan_sqli
        self.scan_headers = scan_headers
        self.scan_redirect = scan_redirect
        self.scan_directories = scan_directories

        self.request_manager = RequestManager(
            timeout=request_timeout,
            delay=request_delay,
        )
        self.classifier = VulnerabilityClassifier()
        self.vulnerabilities: List[Dict[str, Any]] = []
        self.scanned_urls: List[str] = []
        self.scan_start_time: Optional[float] = None
        self.scan_end_time: Optional[float] = None

    def run(
        self,
        progress_callback: Optional[Callable[[str, float, str], None]] = None,
    ) -> List[Dict[str, Any]]:
        """
        Execute the full scanning pipeline.

        Args:
            progress_callback: callable(phase, progress_pct, message)
                - phase: string name of current phase
                - progress_pct: float 0.0 to 1.0
                - message: status message

        Returns:
            List of vulnerability dictionaries with severity scores
        """
        self.scan_start_time = time.time()
        self.vulnerabilities = []

        def _progress(phase, pct, msg):
            if progress_callback:
                progress_callback(phase, pct, msg)
            logger.info(f"[{phase}] {pct*100:.0f}% - {msg}")

        # --- Phase 1: Crawling ---
        _progress("crawl", 0.0, f"Starting crawl of {self.target_url}")

        crawler = WebCrawler(
            base_url=self.target_url,
            request_manager=self.request_manager,
            max_pages=self.max_pages,
        )

        discovered_urls = crawler.crawl(
            progress_callback=lambda curr, total, url: _progress(
                "crawl",
                min(curr / max(total, 1), 0.99),
                f"Crawled: {url[:60]}..."
            )
        )

        self.scanned_urls = discovered_urls
        total_pages = len(discovered_urls)
        _progress("crawl", 1.0, f"Crawl complete. Found {total_pages} pages.")

        # --- Phase 2: Directory Discovery ---
        if self.scan_directories:
            _progress("directories", 0.0, "Probing sensitive directories...")
            dir_vulns = detect_directories(self.target_url, self.request_manager)
            self.vulnerabilities.extend(dir_vulns)
            _progress("directories", 1.0, f"Found {len(dir_vulns)} directory issues.")

        # --- Phase 3: Header Analysis (once on base URL) ---
        if self.scan_headers:
            _progress("headers", 0.0, "Analyzing security headers...")
            header_vulns = detect_missing_headers(self.target_url, self.request_manager)
            self.vulnerabilities.extend(header_vulns)
            _progress("headers", 1.0, f"Found {len(header_vulns)} header issues.")

        # --- Phase 4: Per-page scanning ---
        scan_phases = []
        if self.scan_xss:
            scan_phases.append(("xss", "XSS", detect_xss))
        if self.scan_sqli:
            scan_phases.append(("sqli", "SQL Injection", detect_sqli))
        if self.scan_redirect:
            scan_phases.append(("redirect", "Open Redirect", detect_open_redirect))

        for phase_key, phase_name, detector_fn in scan_phases:
            _progress(phase_key, 0.0, f"Running {phase_name} detection...")
            phase_vulns = []

            # Scan base URL + URLs with params
            urls_to_scan = list({self.target_url} | set(crawler.get_urls_with_params()))
            urls_to_scan = urls_to_scan[:self.max_pages]

            for i, url in enumerate(urls_to_scan):
                pct = (i + 1) / max(len(urls_to_scan), 1)
                _progress(phase_key, pct, f"Testing {phase_name}: {url[:55]}...")

                try:
                    vulns = detector_fn(url, self.request_manager)
                    phase_vulns.extend(vulns)
                except Exception as e:
                    logger.error(f"{phase_name} detector error on {url}: {e}")

            self.vulnerabilities.extend(phase_vulns)
            _progress(phase_key, 1.0, f"{phase_name}: Found {len(phase_vulns)} issues.")

        # --- Phase 5: AI Classification ---
        _progress("ai", 0.0, "Running AI severity classification...")
        classified = self._classify_all(self.vulnerabilities)
        self.vulnerabilities = classified
        _progress("ai", 1.0, f"Classification complete. {len(classified)} total findings.")

        # Deduplicate
        self.vulnerabilities = self._deduplicate(self.vulnerabilities)

        # Sort by severity
        severity_order = {"Critical": 0, "High": 1, "Medium": 2, "Low": 3}
        self.vulnerabilities.sort(
            key=lambda v: severity_order.get(v.get("severity", "Low"), 4)
        )

        self.scan_end_time = time.time()
        return self.vulnerabilities

    def _classify_all(self, vulnerabilities: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """Assign AI-based severity to each vulnerability."""
        for vuln in vulnerabilities:
            severity = self.classifier.classify(vuln)
            vuln["severity"] = severity
            vuln["severity_score"] = self.classifier.get_severity_score(severity)
        return vulnerabilities

    def _deduplicate(self, vulnerabilities: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """Remove duplicate findings based on type + URL + parameter."""
        seen = set()
        unique = []
        for vuln in vulnerabilities:
            key = (vuln.get("type"), vuln.get("url"), vuln.get("parameter"))
            if key not in seen:
                seen.add(key)
                unique.append(vuln)
        return unique

    def get_summary(self) -> Dict[str, Any]:
        """Return a summary of scan results."""
        severity_counts = {"Critical": 0, "High": 0, "Medium": 0, "Low": 0}
        type_counts: Dict[str, int] = {}

        for vuln in self.vulnerabilities:
            sev = vuln.get("severity", "Low")
            severity_counts[sev] = severity_counts.get(sev, 0) + 1
            vtype = vuln.get("type", "Unknown")
            type_counts[vtype] = type_counts.get(vtype, 0) + 1

        elapsed = (
            self.scan_end_time - self.scan_start_time
            if self.scan_start_time and self.scan_end_time
            else 0
        )

        return {
            "target_url": self.target_url,
            "total_vulnerabilities": len(self.vulnerabilities),
            "severity_breakdown": severity_counts,
            "type_breakdown": type_counts,
            "pages_scanned": len(self.scanned_urls),
            "requests_made": self.request_manager.get_request_count(),
            "scan_duration_seconds": round(elapsed, 2),
        }

    def close(self):
        """Clean up resources."""
        self.request_manager.close()