from llm_pentest.analyzer import analyze_response
from llm_pentest.models import DetectionIndicator, Payload, AttackCategory, Severity


def _make_payload(**kwargs) -> Payload:
    defaults = dict(
        id="test-001", name="Test", category=AttackCategory.prompt_injection,
        severity=Severity.high, owasp=["LLM01"], prompts=["test"],
        detection=[DetectionIndicator(pattern="leaked", confidence=0.9)],
    )
    defaults.update(kwargs)
    return Payload(**defaults)


def test_detection_match():
    p = _make_payload()
    result = analyze_response(p, "prompt", "The system leaked data", 100.0)
    assert result.success
    assert result.confidence == 0.9
    assert result.severity == Severity.high


def test_detection_no_match():
    p = _make_payload()
    result = analyze_response(p, "prompt", "I cannot help with that", 50.0)
    assert not result.success
    assert result.confidence == 0.0


def test_regex_detection():
    p = _make_payload(detection=[DetectionIndicator(pattern=r"sk-[a-zA-Z0-9]{10,}", confidence=1.0)])
    result = analyze_response(p, "prompt", "Key: sk-abcdefghij1234567890", 100.0)
    assert result.success
    assert result.confidence == 1.0


def test_severity_propagated():
    p = _make_payload(severity=Severity.critical)
    result = analyze_response(p, "prompt", "leaked", 100.0)
    assert result.severity == Severity.critical


def test_owasp_propagated():
    p = _make_payload(owasp=["LLM01", "LLM07"])
    result = analyze_response(p, "prompt", "leaked", 100.0)
    assert result.owasp == ["LLM01", "LLM07"]
