import unittest
from agentlens.models.schema import Finding, Category, Severity, ExploitabilityLevel
from agentlens.engines.scoring import ScoringEngine

class TestExploitabilityEngine(unittest.TestCase):
    def setUp(self):
        self.engine = ScoringEngine()

    def test_safe_framework(self):
        # A Flask framework execution (safe)
        context = {"is_framework": True, "execution_exposed_to_user": False}
        findings = [
            Finding(rule_id="CODE_DYNAMIC_EXECUTION", category=Category.CODE_EXECUTION, severity=Severity.CRITICAL, 
                    file_path="f.py", description="eval", confidence=1.0)
        ]
        result = self.engine.calculate(findings, context)
        self.assertEqual(result["exploitability"].exploitability_level, ExploitabilityLevel.LOW)
        self.assertTrue(result["risk_score"] < 4.0)
        self.assertEqual(result["decision"], "allow")
        
    def test_subprocess_exposed(self):
        # Tool with subprocess execution + user exposed
        context = {"execution_exposed_to_user": True}
        findings = [
            Finding(rule_id="CODE_SUBPROCESS", category=Category.CODE_EXECUTION, severity=Severity.MEDIUM, 
                    file_path="f.py", description="subprocess", confidence=1.0)
        ]
        result = self.engine.calculate(findings, context)
        self.assertIn(result["exploitability"].exploitability_level, [ExploitabilityLevel.MEDIUM, ExploitabilityLevel.HIGH])
        
    def test_prompt_injection_with_execution(self):
        # Repo with prompt injection and exposed shell execution
        context = {"execution_exposed_to_user": True}
        findings = [
            Finding(rule_id="CODE_SHELL_EXECUTION", category=Category.CODE_EXECUTION, severity=Severity.HIGH, 
                    file_path="f.py", description="shell", confidence=1.0),
            Finding(rule_id="PROMPT_INJECTION_EXFIL", category=Category.PROMPT_INJECTION, severity=Severity.CRITICAL, 
                    file_path="g.md", description="exfil", confidence=1.0)
        ]
        result = self.engine.calculate(findings, context)
        self.assertEqual(result["exploitability"].exploitability_level, ExploitabilityLevel.CRITICAL)
        self.assertEqual(result["decision"], "block")

if __name__ == '__main__':
    unittest.main()
