"""
BREACH.AI - Source Code Analyzer
=================================

White-box analysis for Shannon-style exploitation.
Analyzes JavaScript/TypeScript source code to:
1. Trace data flows from user input to dangerous sinks
2. Identify potentially exploitable paths
3. Generate targeted payloads based on code context

This enables code-aware dynamic testing.
"""

import asyncio
import re
import json
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Set, Tuple
from urllib.parse import urljoin, urlparse

import aiohttp


@dataclass
class DataFlowPath:
    """A potential data flow from source to sink."""
    source: str  # User input source (e.g., req.query.id)
    sink: str  # Dangerous sink (e.g., db.query)
    path: List[str]  # Intermediate steps
    vulnerability_type: str  # sqli, xss, cmdi, etc.
    confidence: float  # 0.0 to 1.0
    code_context: str  # Relevant code snippet
    file_path: str
    line_number: int
    suggested_payloads: List[str] = field(default_factory=list)


@dataclass
class SourceAnalysisResult:
    """Result of source code analysis."""
    target: str
    js_files_analyzed: int
    data_flows_found: List[DataFlowPath]
    endpoints_discovered: List[str]
    parameters_found: List[str]
    technologies_detected: Dict[str, str]
    analysis_time_ms: float


class SourceCodeAnalyzer:
    """
    Analyzes source code for exploitable data flows.

    This implements Shannon's white-box analysis approach:
    1. Fetch and parse JavaScript files
    2. Identify user input sources
    3. Trace data flow to dangerous sinks
    4. Generate targeted exploitation hypotheses
    """

    # User input sources (where attacker-controlled data enters)
    SOURCES = {
        # Express.js
        "req.query": "query_param",
        "req.params": "path_param",
        "req.body": "body_param",
        "req.headers": "header",
        "req.cookies": "cookie",
        "request.query": "query_param",
        "request.body": "body_param",
        # Browser
        "location.hash": "dom_source",
        "location.search": "dom_source",
        "location.href": "dom_source",
        "document.URL": "dom_source",
        "document.referrer": "dom_source",
        "window.name": "dom_source",
        "document.cookie": "cookie",
        # React/Vue/Angular
        "this.props": "props",
        "this.$route.query": "query_param",
        "this.$route.params": "path_param",
        "useSearchParams": "query_param",
        "useParams": "path_param",
        # Generic
        "params": "param",
        "query": "query_param",
        "input": "input",
    }

    # Dangerous sinks (where exploitation occurs)
    SINKS = {
        # SQL Injection
        "query": "sqli",
        ".query(": "sqli",
        ".execute(": "sqli",
        ".raw(": "sqli",
        "sequelize.query": "sqli",
        "knex.raw": "sqli",
        "pool.query": "sqli",
        "db.query": "sqli",
        "$where": "nosql",
        # Command Injection
        "exec(": "cmdi",
        "execSync(": "cmdi",
        "spawn(": "cmdi",
        "spawnSync(": "cmdi",
        "child_process": "cmdi",
        "shell.exec": "cmdi",
        "eval(": "cmdi",
        # XSS
        "innerHTML": "xss",
        "outerHTML": "xss",
        "document.write": "xss",
        "document.writeln": "xss",
        ".html(": "xss",
        "dangerouslySetInnerHTML": "xss",
        "v-html": "xss",
        # Path Traversal
        "readFile": "lfi",
        "readFileSync": "lfi",
        "createReadStream": "lfi",
        "fs.read": "lfi",
        "path.join": "lfi",
        "path.resolve": "lfi",
        # SSRF
        "fetch(": "ssrf",
        "axios(": "ssrf",
        "request(": "ssrf",
        "http.get": "ssrf",
        "https.get": "ssrf",
        "urllib": "ssrf",
        # Template Injection
        "render(": "ssti",
        "compile(": "ssti",
        "template(": "ssti",
        # Deserialization
        "JSON.parse": "deserialization",
        "unserialize": "deserialization",
        "yaml.load": "deserialization",
        # Open Redirect
        "redirect(": "open_redirect",
        "location.href": "open_redirect",
        "window.location": "open_redirect",
    }

    def __init__(
        self,
        session: aiohttp.ClientSession,
        timeout: int = 30,
        max_files: int = 50,
    ):
        self.session = session
        self.timeout = timeout
        self.max_files = max_files

    async def analyze(
        self,
        target: str,
        js_urls: List[str] = None,
        source_code: str = None,
    ) -> SourceAnalysisResult:
        """
        Analyze target for exploitable data flows.

        Args:
            target: Target URL
            js_urls: List of JavaScript file URLs to analyze
            source_code: Optional raw source code to analyze

        Returns:
            SourceAnalysisResult with discovered data flows
        """
        import time
        start_time = time.time()

        data_flows = []
        endpoints = []
        parameters = set()
        technologies = {}
        files_analyzed = 0

        # Analyze provided source code
        if source_code:
            flows = self._analyze_code(source_code, "provided_source", target)
            data_flows.extend(flows)

        # Analyze JavaScript files
        if js_urls:
            for js_url in js_urls[:self.max_files]:
                try:
                    code = await self._fetch_js(js_url)
                    if code:
                        files_analyzed += 1
                        flows = self._analyze_code(code, js_url, target)
                        data_flows.extend(flows)

                        # Extract endpoints and parameters
                        eps, params = self._extract_endpoints(code, target)
                        endpoints.extend(eps)
                        parameters.update(params)

                        # Detect technologies
                        techs = self._detect_technologies(code)
                        technologies.update(techs)

                except Exception as e:
                    pass

        return SourceAnalysisResult(
            target=target,
            js_files_analyzed=files_analyzed,
            data_flows_found=data_flows,
            endpoints_discovered=list(set(endpoints)),
            parameters_found=list(parameters),
            technologies_detected=technologies,
            analysis_time_ms=(time.time() - start_time) * 1000,
        )

    async def _fetch_js(self, url: str) -> Optional[str]:
        """Fetch JavaScript file content."""
        try:
            async with self.session.get(
                url, ssl=False, timeout=self.timeout
            ) as response:
                if response.status == 200:
                    return await response.text()
        except:
            pass
        return None

    def _analyze_code(
        self,
        code: str,
        file_path: str,
        target: str,
    ) -> List[DataFlowPath]:
        """Analyze code for data flow vulnerabilities."""
        data_flows = []

        # Split into lines for context
        lines = code.split("\n")

        for line_num, line in enumerate(lines, 1):
            # Find sources in this line
            for source_pattern, source_type in self.SOURCES.items():
                if source_pattern in line:
                    # Look for sinks in surrounding context
                    context_start = max(0, line_num - 10)
                    context_end = min(len(lines), line_num + 10)
                    context = "\n".join(lines[context_start:context_end])

                    for sink_pattern, vuln_type in self.SINKS.items():
                        if sink_pattern in context:
                            # Found potential data flow
                            flow = self._trace_flow(
                                source=source_pattern,
                                source_type=source_type,
                                sink=sink_pattern,
                                vuln_type=vuln_type,
                                context=context,
                                line=line,
                                line_num=line_num,
                                file_path=file_path,
                            )
                            if flow:
                                data_flows.append(flow)

        return data_flows

    def _trace_flow(
        self,
        source: str,
        source_type: str,
        sink: str,
        vuln_type: str,
        context: str,
        line: str,
        line_num: int,
        file_path: str,
    ) -> Optional[DataFlowPath]:
        """Trace a potential data flow and assess exploitability."""

        # Calculate confidence based on directness of flow
        confidence = 0.5

        # Direct assignment to sink = high confidence
        if source in line and sink in line:
            confidence = 0.9

        # Check for sanitization/validation that would block exploitation
        sanitizers = [
            "escape", "sanitize", "encode", "validate",
            "parseInt", "parseFloat", "Number(",
            "prepared", "parameterized", "placeholder",
            "htmlspecialchars", "strip_tags",
        ]

        for sanitizer in sanitizers:
            if sanitizer.lower() in context.lower():
                confidence *= 0.3  # Significantly reduce confidence

        # Check for common vulnerable patterns that increase confidence
        vulnerable_patterns = [
            (r'\+\s*[\'"]', 0.2),  # String concatenation
            (r'`\$\{', 0.2),  # Template literal injection
            (r'\.query\s*\(.*\+', 0.3),  # SQL string concatenation
            (r'innerHTML\s*=', 0.2),  # Direct innerHTML assignment
            (r'eval\s*\(', 0.3),  # Eval usage
        ]

        for pattern, boost in vulnerable_patterns:
            if re.search(pattern, context):
                confidence = min(1.0, confidence + boost)

        # Don't report very low confidence flows
        if confidence < 0.3:
            return None

        # Generate suggested payloads based on vulnerability type
        payloads = self._generate_payloads(vuln_type, source_type)

        return DataFlowPath(
            source=source,
            sink=sink,
            path=[source, "...", sink],
            vulnerability_type=vuln_type,
            confidence=confidence,
            code_context=context[:500],
            file_path=file_path,
            line_number=line_num,
            suggested_payloads=payloads,
        )

    def _generate_payloads(
        self,
        vuln_type: str,
        source_type: str,
    ) -> List[str]:
        """Generate payloads based on vulnerability type."""
        payloads = {
            "sqli": [
                "' OR '1'='1",
                "' UNION SELECT NULL--",
                "'; DROP TABLE users--",
                "' AND SLEEP(5)--",
            ],
            "xss": [
                "<script>alert('XSS')</script>",
                "<img src=x onerror=alert('XSS')>",
                "javascript:alert('XSS')",
                "<svg onload=alert('XSS')>",
            ],
            "cmdi": [
                "; id",
                "| whoami",
                "`id`",
                "$(whoami)",
            ],
            "lfi": [
                "../../../etc/passwd",
                "....//....//etc/passwd",
                "/etc/passwd%00",
                "..\\..\\..\\windows\\win.ini",
            ],
            "ssrf": [
                "http://127.0.0.1:80",
                "http://localhost:22",
                "http://169.254.169.254/latest/meta-data/",
                "file:///etc/passwd",
            ],
            "ssti": [
                "{{7*7}}",
                "${7*7}",
                "#{7*7}",
                "<%= 7*7 %>",
            ],
            "nosql": [
                '{"$ne": null}',
                '{"$gt": ""}',
                '{"$regex": ".*"}',
                '{"$where": "1==1"}',
            ],
            "open_redirect": [
                "//evil.com",
                "https://evil.com",
                "/\\evil.com",
                "//evil.com/%2f..",
            ],
        }

        return payloads.get(vuln_type, [])

    def _extract_endpoints(
        self,
        code: str,
        base_url: str,
    ) -> Tuple[List[str], Set[str]]:
        """Extract API endpoints and parameters from code."""
        endpoints = []
        parameters = set()

        # API endpoint patterns
        endpoint_patterns = [
            # Express.js routes
            r'(?:app|router)\.(get|post|put|delete|patch)\s*\([\'"]([^"\']+)[\'"]',
            # Fetch/axios calls
            r'(?:fetch|axios\.get|axios\.post)\s*\([\'"]([^"\']+)[\'"]',
            # API paths in strings
            r'[\'"]\/api\/[^\'"]+[\'"]',
            r'[\'"]\/v\d+\/[^\'"]+[\'"]',
        ]

        for pattern in endpoint_patterns:
            matches = re.findall(pattern, code)
            for match in matches:
                if isinstance(match, tuple):
                    endpoint = match[-1]  # Last group is the path
                else:
                    endpoint = match

                if endpoint.startswith("/"):
                    endpoints.append(urljoin(base_url, endpoint))
                elif endpoint.startswith("http"):
                    endpoints.append(endpoint)

        # Parameter patterns
        param_patterns = [
            r'req\.query\.(\w+)',
            r'req\.params\.(\w+)',
            r'req\.body\.(\w+)',
            r'params\[[\'"]([\w]+)[\'"]\]',
            r'searchParams\.get\([\'"](\w+)[\'"]\)',
            r'name=[\'"](\w+)[\'"]',
        ]

        for pattern in param_patterns:
            matches = re.findall(pattern, code)
            parameters.update(matches)

        return endpoints, parameters

    def _detect_technologies(self, code: str) -> Dict[str, str]:
        """Detect technologies/frameworks from code."""
        technologies = {}

        tech_patterns = {
            "react": [r'import.*from\s+[\'"]react[\'"]', r'React\.'],
            "vue": [r'import.*from\s+[\'"]vue[\'"]', r'Vue\.'],
            "angular": [r'@angular/', r'@Component'],
            "express": [r'express\(\)', r'app\.listen'],
            "nextjs": [r'next/router', r'getServerSideProps'],
            "jquery": [r'\$\(', r'jQuery'],
            "axios": [r'axios\.', r'import axios'],
            "graphql": [r'graphql', r'gql`'],
            "firebase": [r'firebase/', r'firebaseConfig'],
            "supabase": [r'@supabase/', r'supabaseUrl'],
            "prisma": [r'@prisma/', r'prisma\.'],
            "sequelize": [r'sequelize', r'Sequelize'],
        }

        for tech, patterns in tech_patterns.items():
            for pattern in patterns:
                if re.search(pattern, code, re.IGNORECASE):
                    technologies[tech] = "detected"
                    break

        return technologies


class HypothesisGenerator:
    """
    Generates exploitation hypotheses from data flow analysis.

    Based on Shannon's approach of generating "hypothesized exploitable paths"
    that are then validated through actual exploitation.
    """

    @dataclass
    class ExploitHypothesis:
        """A hypothesis about an exploitable vulnerability."""
        target: str
        parameter: str
        vulnerability_type: str
        confidence: float
        payloads: List[str]
        reasoning: str
        data_flow: Optional[DataFlowPath] = None

    def generate(
        self,
        analysis_result: SourceAnalysisResult,
    ) -> List[ExploitHypothesis]:
        """Generate exploitation hypotheses from analysis."""
        hypotheses = []

        for flow in analysis_result.data_flows_found:
            hypothesis = self.ExploitHypothesis(
                target=analysis_result.target,
                parameter=self._extract_param_name(flow.source),
                vulnerability_type=flow.vulnerability_type,
                confidence=flow.confidence,
                payloads=flow.suggested_payloads,
                reasoning=f"Data flow from {flow.source} to {flow.sink} detected in {flow.file_path}",
                data_flow=flow,
            )
            hypotheses.append(hypothesis)

        # Sort by confidence
        hypotheses.sort(key=lambda h: h.confidence, reverse=True)

        return hypotheses

    def _extract_param_name(self, source: str) -> str:
        """Extract parameter name from source identifier."""
        # req.query.id -> id
        parts = source.split(".")
        if len(parts) > 2:
            return parts[-1]
        return source
