#!/usr/bin/env python3

"""
Project: BRS-XSS (XSS Detection Suite)
Company: EasyProTech LLC (www.easypro.tech)
Dev: Brabus
Date: Sat 26 Oct 2025 13:43:19 UTC
Status: Modified
Telegram: https://t.me/EasyProTech
"""

import asyncio
from typing import Optional
from urllib.parse import urlparse, urljoin

import typer
from rich.console import Console
from rich.progress import Progress
from rich.table import Table

from brsxss import VERSION
from brsxss.core.scanner import XSSScanner
from brsxss.report.report_generator import ReportGenerator
from brsxss.report.report_types import ReportConfig, ReportFormat
from brsxss.report.data_models import VulnerabilityData, ScanStatistics
from brsxss.utils.logger import Logger
from brsxss.crawler.data_models import DiscoveredParameter

console = Console()


async def simple_scan(
    target: str,
    threads: int = 10,
    timeout: int = 15,
    output: Optional[str] = None,
    deep: bool = False,
    verbose: bool = False,
    ml_mode: bool = False,
    blind_xss_webhook: Optional[str] = None,
    no_ssl_verify: bool = False,
    safe_mode: bool = True,
    pool_cap: int = 10000,
    max_payloads: int = 500,
):
    """Scan target for XSS vulnerabilities - specify domain or IP only"""
    
    # Configure logging based on verbosity
    Logger.configure_logging(verbose=verbose)

    logger = Logger("cli.simple_scan")
    
    console.print(f"[bold green]BRS-XSS v{VERSION}[/bold green]")
    console.print(f"Target: {target}")
    
    if verbose:
        console.print("[dim]Verbose mode enabled - detailed parameter analysis[/dim]")
    if ml_mode:
        console.print("[dim]ML mode enabled - vulnerability classification[/dim]")
    
    # Scanner will be created with progress callback during scanning
    
    if blind_xss_webhook:
        console.print(f"Blind XSS webhook enabled: {blind_xss_webhook}")
    
    try:
        # Auto-detect protocol and build URLs
        # Force HTTP for internal IPs when SSL verification is disabled
        force_http = no_ssl_verify and (
            target.startswith('192.168.') or target.startswith('10.') or target.startswith('172.') or
            target.startswith('127.') or 'localhost' in target
        )
        scan_targets = _build_scan_targets(target, force_http)
        
        console.print(f"Auto-detected {len(scan_targets)} targets to scan")
        
        all_vulnerabilities = []
        
        with Progress() as progress:
            # Create main task for URLs
            url_task = progress.add_task("Scanning targets...", total=len(scan_targets))
            # Create detailed task for payloads
            payload_task = progress.add_task("Testing payloads...", total=100)
            
            def update_payload_progress(current: int, total: int):
                """Update payload progress bar"""
                if total > 0:
                    percentage = min(100, (current * 100) // total)
                    progress.update(payload_task, completed=percentage, total=100,
                                  description=f"Testing payload {current}/{total}")
            
            # Create a single reusable HTTP client
            from brsxss.core.http_client import HTTPClient
            http_client = HTTPClient(timeout=timeout, verify_ssl=not no_ssl_verify)
            
            try:
                for url in scan_targets:
                    progress.update(url_task, description=f"Scanning {url}")
                    progress.update(payload_task, completed=0, description="Discovering parameters...")
                    
                    try:
                        # Auto-discover parameters (entry points)
                        discovered_params = await _discover_parameters(url, deep, http_client)
                        
                        if discovered_params:
                            progress.print(f"Found {len(discovered_params)} entry points in {url}")
                            
                            # Create scanner with progress callback
                            scanner_with_progress = XSSScanner(
                                timeout=timeout, 
                                max_concurrent=threads, 
                                verify_ssl=not no_ssl_verify, 
                                blind_xss_webhook=blind_xss_webhook,
                                progress_callback=update_payload_progress,
                                http_client=http_client  # Reuse client
                            )
                            
                            # Scan each discovered entry point
                            for entry_point in discovered_params:
                                vulns = await scanner_with_progress.scan_url(
                                    entry_point.url, 
                                    entry_point.method,
                                    entry_point.params
                                )
                                all_vulnerabilities.extend(vulns)
                            
                        progress.advance(url_task)
                        progress.update(payload_task, completed=100, description="URL scan completed")
                        
                    except Exception as e:
                        logger.warning(f"Error scanning {url}: {e}")
                        progress.advance(url_task)
            
            finally:
                # Ensure the shared HTTP client is closed
                await http_client.close()

        # Display results in a rich table
        console.print("\n[bold]Scan Summary[/bold]")
        
        if all_vulnerabilities:
            def _sev_str(value):
                try:
                    # Enum with value
                    v = getattr(value, 'value', None)
                    if isinstance(v, str):
                        return v.lower()
                    # Enum name
                    n = getattr(value, 'name', None)
                    if isinstance(n, str):
                        return n.lower()
                except Exception:
                    pass
                if isinstance(value, str):
                    return value.lower()
                return 'low'

            table = Table(title="Vulnerabilities Found", show_header=True, header_style="bold magenta")
            table.add_column("Severity", style="dim", width=12)
            table.add_column("URL", style="cyan")
            table.add_column("Method", style="green")
            table.add_column("Parameter", style="yellow")
            table.add_column("Payload Snippet", style="white")

            severity_colors = {
                "critical": "bold red",
                "high": "red",
                "medium": "yellow",
                "low": "cyan",
                "info": "green"
            }

            # Sort vulnerabilities by severity (high to low)
            severities = ["critical", "high", "medium", "low", "info"]
            sorted_vulns = sorted(
                all_vulnerabilities,
                key=lambda v: severities.index(_sev_str(v.get('severity', 'low'))) if _sev_str(v.get('severity', 'low')) in severities else 99
            )

            for vuln in sorted_vulns:
                severity = _sev_str(vuln.get('severity', 'low'))
                color = severity_colors.get(severity, "white")
                
                payload = vuln.get('payload', '')
                payload_snippet = (payload[:40] + '...') if len(payload) > 40 else payload

                table.add_row(
                    f"[{color}]{severity.upper()}[/{color}]",
                    vuln.get('url', ''),
                    vuln.get('http_method', 'GET'),
                    vuln.get('parameter', ''),
                    payload_snippet
                )
            
            console.print(table)
        else:
            console.print("[green]✔ No vulnerabilities found - target appears secure.[/green]")
        
        # Save report
        if not output:
            # Default output path (sanitized + ensured directories)
            import os
            from brsxss.utils.paths import sanitize_filename, ensure_dir
            timestamp = int(__import__('time').time())
            clean_target = sanitize_filename(target, max_len=50)
            filename = f"scan_report_{clean_target}_{timestamp}.json"
            results_dir = os.path.abspath("results/json")
            ensure_dir(results_dir)
            output = os.path.join(results_dir, filename)
        
        _save_simple_report(all_vulnerabilities, scan_targets, output)
        console.print(f"Report saved: {output}")

        # Generate multi-format report (HTML + JSON)
        try:
            # Convert to VulnerabilityData
            vuln_items = []
            for idx, v in enumerate(all_vulnerabilities, 1):
                severity = v.get('severity')
                if hasattr(severity, 'value'):
                    severity = severity.value  # type: ignore[union-attr]
                elif not isinstance(severity, str):
                    severity = 'low'
                vuln_items.append(
                    VulnerabilityData(
                        id=f"xss_{idx}",
                        title=f"XSS in parameter {v.get('parameter','')}",
                        description=f"Possible XSS detected for parameter {v.get('parameter','')}.",
                        severity=severity,
                        confidence=float(v.get('confidence', 0.5)),
                        url=v.get('url',''),
                        parameter=v.get('parameter',''),
                        payload=v.get('payload',''),
                    )
                )

            # Scan statistics
            stats = ScanStatistics(
                total_urls_tested=len(scan_targets),
                total_parameters_tested=sum(len((await _discover_parameters(u, False, None))) for u in scan_targets if isinstance(u,str)) if False else 0,
                total_payloads_tested=0,
                total_requests_sent=0,
                scan_duration=0.0,
                vulnerabilities_found=len(vuln_items),
            )

            # Configure report
            report_config = ReportConfig(
                title=f"BRS-XSS Scan Report - {scan_targets[0] if scan_targets else target}",
                output_dir="results",
                filename_template="brsxss_report_{timestamp}",
                formats=[ReportFormat.HTML, ReportFormat.JSON],
                include_recommendations=True,
                include_methodology=True,
            )
            generator = ReportGenerator(report_config)
            # Load config for report generation
            from brsxss.core.config_manager import ConfigManager
            config_mgr = ConfigManager()
            
            policy = {
                'min_vulnerability_score': config_mgr.get('scanner.min_vulnerability_score', 2.0),
                'severity_bands': {
                    'critical': '>= 9.0', 'high': '>=7.0', 'medium': '>=4.0', 'low': '>=1.0', 'info': '>0'
                }
            }
            generated = generator.generate_report(vuln_items, stats, target_info={"url": target, "policy": policy})

            # Move reports to structured directories
            import os
            os.makedirs("results/html", exist_ok=True)
            os.makedirs("results/json", exist_ok=True)
            for fmt, path in generated.items():
                try:
                    if path.endswith('.html'):
                        new_path = os.path.join("results/html", os.path.basename(path))
                    elif path.endswith('.json'):
                        new_path = os.path.join("results/json", os.path.basename(path))
                    else:
                        new_path = path
                    if new_path != path:
                        os.replace(path, new_path)
                        path = new_path
                except Exception as move_err:
                    logger.debug(f"Report move error: {move_err}")
                console.print(f"Report generated: {path}")
        except Exception as e:
            logger.debug(f"Failed to generate report: {e}")
    
    except Exception as e:
        logger.error(f"Scan failed: {e}")
        raise typer.Exit(1)
    
    finally:
        # Clean up HTTP sessions
        try:
            # Give time for pending requests to complete
            await asyncio.sleep(0.5)
            # Note: Individual scanners are cleaned up in the loop above
            # Additional delay to ensure SSL cleanup
            await asyncio.sleep(0.5)
        except Exception as e:
            logger.debug(f"Error in cleanup: {e}")


def _build_scan_targets(target: str, force_http: bool = False) -> list:
    """Build list of URLs to scan from target domain/IP"""
    
    # Clean target
    target = target.strip()
    
    # Check if target is already a full URL with path/query
    if target.startswith(('http://', 'https://')):
        # User provided full URL - use it directly
        return [target]
    elif '/' in target or '?' in target:
        # User provided domain with path/query - add protocols
        # If non-SSL port is explicitly specified, avoid adding https
        non_ssl_port = (":" in target) and not (":443" in target or ":8443" in target)
        if force_http or non_ssl_port:
            return [f"http://{target}"]
        else:
            return [f"http://{target}", f"https://{target}"]
    
    # User provided only domain/IP - generate common endpoints
    target = target.lower()
    
    # Build target URLs
    targets = []
    
    # Smart protocol selection
    if force_http:
        # Force HTTP only for internal IPs or when SSL issues
        base_urls = [f"http://{target}"]
    else:
        # Try both HTTP and HTTPS for external domains
        base_urls = [f"http://{target}", f"https://{target}"]
    
    # Common paths to test
    common_paths = [
        "/",
        "/index.php",
        "/search.php", 
        "/login.php",
        "/contact.php",
        "/search",
        "/api/search",
        "/search?q=test",
        "/index.php?page=test",
        "/search.php?search=test",
        "/contact.php?name=test&email=test"
    ]
    
    for base_url in base_urls:
        for path in common_paths:
            targets.append(urljoin(base_url, path))
    
    return targets


async def _discover_parameters(url: str, deep_scan: bool = False, http_client=None) -> list[DiscoveredParameter]:
    """Auto-discover parameters in URL and forms, returning structured entry points."""
    
    entry_points = []
    
    # 1. Extract parameters from the initial URL itself (GET request)
    from urllib.parse import parse_qs
    parsed = urlparse(url)
    url_params = parse_qs(parsed.query)
    
    if url_params:
        # Create an entry point for GET parameters in the URL
        # Keep the full URL with query string - scanner needs it
        entry_points.append(
            DiscoveredParameter(
                url=url,
                method="GET",
                params={param: values[0] if values else "test" for param, values in url_params.items()}
            )
        )

    # 2. If deep scan, crawl and extract forms
    if deep_scan and http_client:
        try:
            from brsxss.crawler.engine import CrawlerEngine, CrawlConfig
            from brsxss.crawler.form_extractor import FormExtractor
            
            config = CrawlConfig(
                max_depth=2, 
                max_urls=20, 
                max_concurrent=5, # Increased concurrency
                timeout=15,
                extract_forms=True,
                extract_links=True
            )
            crawler = CrawlerEngine(config, http_client)
            crawl_results = await crawler.crawl(url)
            
            form_extractor = FormExtractor()
            
            for result in crawl_results:
                if result.status_code == 200 and result.content:
                    forms = form_extractor.extract_forms(result.content, result.url)
                    
                    for form in forms:
                        # Create an entry point for each discovered form
                        form_params = {}
                        for field in form.testable_fields:
                            if field.field_type.name == 'PASSWORD':
                                form_params[field.name] = "password123"
                            elif field.field_type.name == 'EMAIL':
                                form_params[field.name] = "test@example.com"
                            else:
                                form_params[field.name] = "test"
                        
                        entry_points.append(
                            DiscoveredParameter(
                                url=form.action,
                                method=form.method.upper(),
                                params=form_params
                            )
                        )

        except Exception as e:
            logger = Logger("cli.simple_scan._discover_parameters")
            logger.error(f"Deep discovery failed: {e}")
            # Fallback to basic regex on the initial page if deep scan fails
            try:
                response = await http_client.get(url)
                if response.status_code == 200:
                    import re
                    form_inputs = re.findall(r'<input[^>]*name=["\']([^"\']+)["\']', response.text, re.I)
                    if form_inputs:
                        entry_points.append(
                            DiscoveredParameter(
                                url=url,
                                method="GET", # Assume GET as a fallback
                                params={name: "test" for name in form_inputs}
                            )
                        )
            except Exception:
                pass # Ignore fallback errors
    
    # Remove duplicate entry points
    unique_entry_points = []
    seen = set()
    for ep in entry_points:
        # Create a unique key for each entry point
        key = (ep.url, ep.method, tuple(sorted(ep.params.keys())))
        if key not in seen:
            unique_entry_points.append(ep)
            seen.add(key)
            
    return unique_entry_points


def _save_simple_report(vulnerabilities: list, targets: list, output_path: str):
    """Save simple scan report"""
    
    import json
    from datetime import datetime
    from enum import Enum
    
    # Custom JSON encoder for Enum types
    class CustomJSONEncoder(json.JSONEncoder):
        def default(self, obj):
            if isinstance(obj, Enum):
                return obj.value
            return super().default(obj)
    
    # Convert vulnerabilities to serializable format recursively
    def make_serializable(obj):
        from dataclasses import is_dataclass, asdict
        
        if isinstance(obj, Enum):
            return obj.value
        elif is_dataclass(obj):
            return {k: make_serializable(v) for k, v in asdict(obj).items()}
        elif isinstance(obj, dict):
            return {k: make_serializable(v) for k, v in obj.items()}
        elif isinstance(obj, (list, tuple)):
            return [make_serializable(item) for item in obj]
        else:
            return obj
    
    serializable_vulns = [make_serializable(vuln) for vuln in vulnerabilities]
    
    report = {
        "scan_info": {
            "timestamp": datetime.now().isoformat(),
            "scanner": f"BRS-XSS Simple Scanner v{VERSION}",
            "targets_scanned": len(targets),
            "vulnerabilities_found": len(serializable_vulns)
        },
        "targets": targets,
        "vulnerabilities": serializable_vulns
    }
    
    # Determine format by extension with atomic write
    from brsxss.utils.paths import atomic_write
    content = json.dumps(report, indent=2, cls=CustomJSONEncoder)
    if output_path.endswith('.json'):
        atomic_write(output_path, content)
    else:
        atomic_write(output_path + '.json', content)


def simple_scan_wrapper(
    target: str,
    threads: int = 10,
    timeout: int = 15,
    output: Optional[str] = None,
    deep: bool = False,
    verbose: bool = False,
    ml_mode: bool = False,
    blind_xss_webhook: Optional[str] = None,
    no_ssl_verify: bool = False,
    safe_mode: bool = True,
    pool_cap: int = 10000,
    max_payloads: int = 500,
):
    """Wrapper to run async scan function"""
    return asyncio.run(simple_scan(target, threads, timeout, output, deep, verbose, ml_mode, blind_xss_webhook, no_ssl_verify, safe_mode, pool_cap, max_payloads))

# Create typer app for this command
app = typer.Typer()
app.command()(simple_scan_wrapper)