"""PRIV-001: Privilege Escalation / Excessive Agency detection rule.

Detects patterns where tool functions have overly broad capabilities
that could be exploited for privilege escalation, such as shell execution,
eval(), dynamic imports, or unrestricted filesystem access.

NOTE: This module contains string references to unsafe functions because it is
a SECURITY SCANNER that detects these patterns in target code. These are
detection signatures, not actual usage.
"""

from __future__ import annotations

import ast

from . import Rule
from ..models import Finding, Severity

# Dangerous execution patterns — function names to detect
SHELL_EXEC_FUNCS = {"os.system", "os.popen"}

# Subprocess with shell=True
SUBPROCESS_FUNCS = {
    "subprocess.run", "subprocess.call", "subprocess.Popen",
    "subprocess.check_output", "subprocess.check_call",
}

# Dangerous built-in functions
DANGEROUS_BUILTINS = {"eval", "exec"}

# Dynamic import function
DYNAMIC_IMPORT = {"__import__"}

# Unrestricted filesystem patterns
UNRESTRICTED_FS_FUNCS = {"shutil.rmtree", "os.remove", "os.unlink", "os.rmdir"}


class PrivilegeEscalationRule(Rule):
    """Detects excessive agency / privilege escalation in tool handlers.

    Flags tool functions that use overly broad capabilities such as
    os.system(), subprocess with shell=True, eval(), __import__(),
    or unrestricted filesystem deletion.
    """

    @property
    def rule_id(self) -> str:
        return "PRIV-001"

    @property
    def name(self) -> str:
        return "Privilege Escalation / Excessive Agency"

    @property
    def severity(self) -> Severity:
        return Severity.HIGH

    @property
    def cwe(self) -> str:
        return "CWE-269"

    @property
    def description(self) -> str:
        return (
            "Tool handler uses overly broad capabilities that could "
            "enable privilege escalation."
        )

    @property
    def detailed_description(self) -> str:
        return (
            "Tool handler uses overly broad capabilities that could "
            "enable privilege escalation.\n\n"
            "MCP tool handlers should follow the principle of least privilege.\n"
            "Functions that can run arbitrary commands, evaluate code,\n"
            "or perform unrestricted filesystem operations give excessive\n"
            "agency to the AI agent, enabling privilege escalation.\n\n"
            "How to fix:\n"
            "  - Replace os.system()/shell=True with subprocess argument lists\n"
            "  - Remove eval() and use safe alternatives\n"
            "  - Replace __import__() with explicit imports\n"
            "  - Restrict filesystem operations to specific directories"
        )

    def check(self, tree: ast.Module, source: str) -> list[Finding]:
        """Check for excessive agency patterns in tool handlers."""
        findings: list[Finding] = []
        for node in ast.walk(tree):
            if not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
                continue
            if not self._is_tool_handler(node):
                continue
            findings.extend(self._check_shell_exec(node, source))
            findings.extend(self._check_shell_subprocess(node, source))
            findings.extend(self._check_dangerous_builtins(node, source))
            findings.extend(self._check_dynamic_import(node, source))
            findings.extend(self._check_unrestricted_fs(node, source))
        return findings

    def _check_shell_exec(
        self,
        func_node: ast.FunctionDef | ast.AsyncFunctionDef,
        source: str,
    ) -> list[Finding]:
        """Detect os.system() and os.popen() calls."""
        findings: list[Finding] = []
        for child in ast.walk(func_node):
            if not isinstance(child, ast.Call):
                continue
            func_name = self._get_func_name(child)
            if func_name in SHELL_EXEC_FUNCS:
                findings.append(Finding(
                    rule_id=self.rule_id,
                    severity=self.severity,
                    title=f"Excessive agency: {func_name}() in tool handler",
                    description=(
                        f"Tool handler '{func_node.name}' uses "
                        f"{func_name}() which runs arbitrary commands."
                    ),
                    file="",
                    line=child.lineno,
                    code_snippet=self._get_snippet(source, child.lineno),
                    recommendation=(
                        "Use subprocess with argument list (no shell=True) "
                        "and restrict allowed commands."
                    ),
                    cwe=self.cwe,
                    confidence="HIGH",
                ))
        return findings

    def _check_shell_subprocess(
        self,
        func_node: ast.FunctionDef | ast.AsyncFunctionDef,
        source: str,
    ) -> list[Finding]:
        """Detect subprocess calls with shell=True."""
        findings: list[Finding] = []
        for child in ast.walk(func_node):
            if not isinstance(child, ast.Call):
                continue
            func_name = self._get_func_name(child)
            if func_name not in SUBPROCESS_FUNCS:
                continue
            shell_true = any(
                isinstance(kw.value, ast.Constant) and kw.value.value is True
                for kw in child.keywords
                if kw.arg == "shell"
            )
            if shell_true:
                findings.append(Finding(
                    rule_id=self.rule_id,
                    severity=self.severity,
                    title="Excessive agency: subprocess shell=True in tool",
                    description=(
                        f"Tool handler '{func_node.name}' uses "
                        f"{func_name}(shell=True) enabling arbitrary commands."
                    ),
                    file="",
                    line=child.lineno,
                    code_snippet=self._get_snippet(source, child.lineno),
                    recommendation=(
                        "Remove shell=True and pass commands as a list. "
                        "Restrict to specific allowed commands."
                    ),
                    cwe=self.cwe,
                    confidence="HIGH",
                ))
        return findings

    def _check_dangerous_builtins(
        self,
        func_node: ast.FunctionDef | ast.AsyncFunctionDef,
        source: str,
    ) -> list[Finding]:
        """Detect eval() calls in tool handlers."""
        findings: list[Finding] = []
        for child in ast.walk(func_node):
            if not isinstance(child, ast.Call):
                continue
            if isinstance(child.func, ast.Name) and child.func.id in DANGEROUS_BUILTINS:
                builtin_name = child.func.id
                findings.append(Finding(
                    rule_id=self.rule_id,
                    severity=Severity.CRITICAL,
                    title=f"Excessive agency: {builtin_name}() in tool handler",
                    description=(
                        f"Tool handler '{func_node.name}' uses "
                        f"{builtin_name}() which can run arbitrary code."
                    ),
                    file="",
                    line=child.lineno,
                    code_snippet=self._get_snippet(source, child.lineno),
                    recommendation=(
                        f"Remove {builtin_name}() and use safe alternatives. "
                        f"Never evaluate user-controlled strings as code."
                    ),
                    cwe="CWE-95",
                    confidence="HIGH",
                ))
        return findings

    def _check_dynamic_import(
        self,
        func_node: ast.FunctionDef | ast.AsyncFunctionDef,
        source: str,
    ) -> list[Finding]:
        """Detect __import__() calls."""
        findings: list[Finding] = []
        for child in ast.walk(func_node):
            if not isinstance(child, ast.Call):
                continue
            if isinstance(child.func, ast.Name) and child.func.id == "__import__":
                findings.append(Finding(
                    rule_id=self.rule_id,
                    severity=self.severity,
                    title="Excessive agency: __import__() in tool handler",
                    description=(
                        f"Tool handler '{func_node.name}' uses "
                        f"__import__() for dynamic imports."
                    ),
                    file="",
                    line=child.lineno,
                    code_snippet=self._get_snippet(source, child.lineno),
                    recommendation=(
                        "Use explicit imports instead of __import__(). "
                        "Dynamic imports can load arbitrary modules."
                    ),
                    cwe=self.cwe,
                    confidence="HIGH",
                ))
        return findings

    def _check_unrestricted_fs(
        self,
        func_node: ast.FunctionDef | ast.AsyncFunctionDef,
        source: str,
    ) -> list[Finding]:
        """Detect unrestricted filesystem deletion operations."""
        findings: list[Finding] = []
        for child in ast.walk(func_node):
            if not isinstance(child, ast.Call):
                continue
            func_name = self._get_func_name(child)
            if func_name in UNRESTRICTED_FS_FUNCS:
                findings.append(Finding(
                    rule_id=self.rule_id,
                    severity=self.severity,
                    title=f"Excessive agency: {func_name}() in tool handler",
                    description=(
                        f"Tool handler '{func_node.name}' uses "
                        f"{func_name}() without path restrictions."
                    ),
                    file="",
                    line=child.lineno,
                    code_snippet=self._get_snippet(source, child.lineno),
                    recommendation=(
                        "Restrict filesystem operations to specific "
                        "directories. Validate paths before deletion."
                    ),
                    cwe=self.cwe,
                    confidence="MEDIUM",
                ))
        return findings

    def _is_tool_handler(
        self, node: ast.FunctionDef | ast.AsyncFunctionDef
    ) -> bool:
        """Check if function is an MCP tool handler."""
        name_lower = node.name.lower()
        if any(kw in name_lower for kw in ("tool", "handle")):
            return True
        for dec in node.decorator_list:
            if "tool" in ast.dump(dec).lower():
                return True
        return False
