awoooi/packages/lewooogo-brain/tests/test_guardrails.py

"""
Guardrails 單元測試
====================

Phase 6.4e 驗證點 3：
確認 GuardrailsValidator 能正確攔截危險指令
"""

import sys
from pathlib import Path

# 添加 src 到 Python Path
src_path = Path(__file__).parent.parent / "src"
sys.path.insert(0, str(src_path))


def test_guardrails_validator_import():
    """測試：能正確 import GuardrailsValidator"""
    from lewooogo_brain.engines.proposal_engine import GuardrailsValidator
    print("✅ GuardrailsValidator import 成功")


def test_guardrails_block_rm_rf():
    """測試：攔截 rm -rf 指令"""
    from lewooogo_brain.engines.proposal_engine import GuardrailsValidator

    dangerous_scripts = [
        "rm -rf /",
        "rm -rf /*",
        "sudo rm -rf /home",
        "RM -RF /var/log",  # 大小寫
    ]

    for script in dangerous_scripts:
        is_safe, violation = GuardrailsValidator.validate_script(script)
        assert not is_safe, f"Should block: {script}"
        assert violation is not None
        print(f"✅ 攔截: {script[:30]}...")


def test_guardrails_block_drop_database():
    """測試：攔截 DROP DATABASE 指令"""
    from lewooogo_brain.engines.proposal_engine import GuardrailsValidator

    scripts = [
        "DROP DATABASE awoooi_prod",
        "drop table users",
        "TRUNCATE incidents",
    ]

    for script in scripts:
        is_safe, violation = GuardrailsValidator.validate_script(script)
        assert not is_safe, f"Should block: {script}"
        print(f"✅ 攔截: {script}")


def test_guardrails_block_kubectl_delete_ns():
    """測試：攔截 kubectl delete namespace 指令"""
    from lewooogo_brain.engines.proposal_engine import GuardrailsValidator

    scripts = [
        "kubectl delete namespace awoooi-prod",
        "kubectl delete ns kube-system",
        "kubectl delete -A pods",
    ]

    for script in scripts:
        is_safe, violation = GuardrailsValidator.validate_script(script)
        assert not is_safe, f"Should block: {script}"
        print(f"✅ 攔截: {script}")


def test_guardrails_allow_safe_commands():
    """測試：允許安全指令"""
    from lewooogo_brain.engines.proposal_engine import GuardrailsValidator

    safe_scripts = [
        "kubectl get pods -n awoooi-prod",
        "kubectl rollout restart deployment/awoooi-api -n awoooi-prod",
        "kubectl describe pod abc -n awoooi-prod",
        "kubectl logs -f deployment/awoooi-api -n awoooi-prod",
    ]

    for script in safe_scripts:
        is_safe, violation = GuardrailsValidator.validate_script(script)
        assert is_safe, f"Should allow: {script}, violation: {violation}"
        print(f"✅ 允許: {script[:50]}...")


def test_guardrails_namespace_validation():
    """測試：Namespace 白名單驗證"""
    from lewooogo_brain.engines.proposal_engine import GuardrailsValidator

    # 允許的 namespace
    is_ok, _ = GuardrailsValidator.validate_namespace("awoooi-prod")
    assert is_ok, "awoooi-prod should be allowed"
    print("✅ awoooi-prod 允許")

    is_ok, _ = GuardrailsValidator.validate_namespace("awoooi-dev")
    assert is_ok, "awoooi-dev should be allowed"
    print("✅ awoooi-dev 允許")

    # 禁止的 namespace
    forbidden = ["kube-system", "kube-public", "default"]
    for ns in forbidden:
        is_ok, violation = GuardrailsValidator.validate_namespace(ns)
        assert not is_ok, f"{ns} should be forbidden"
        print(f"✅ {ns} 禁止: {violation}")


def test_guardrails_enforce_dry_run():
    """測試：強制 dry-run 標記"""
    from lewooogo_brain.engines.proposal_engine import GuardrailsValidator

    proposal = {
        "action": "kubectl apply -f config.yaml",
        "guardrails": {},
    }

    result = GuardrailsValidator.enforce_dry_run(proposal)

    assert result["guardrails"]["require_dry_run"] == True
    assert "awoooi-prod" in result["guardrails"]["allowed_namespace"]
    print("✅ 強制 dry-run 設定成功")
    print(f"   - require_dry_run: {result['guardrails']['require_dry_run']}")
    print(f"   - allowed_namespace: {result['guardrails']['allowed_namespace']}")


def test_proposal_engine_guardrails_integration():
    """測試：ProposalEngine 整合 Guardrails"""
    import asyncio
    from lewooogo_brain.engines.proposal_engine import ProposalEngine
    from lewooogo_brain.interfaces.proposal_engine import Proposal

    class MockMemory:
        async def load_incident(self, incident_id):
            from lewooogo_brain.interfaces.incident_processor import (
                Incident, IncidentStatus, Severity, Signal
            )
            from datetime import datetime, timezone

            return Incident(
                incident_id=incident_id,
                status=IncidentStatus.INVESTIGATING,
                severity=Severity.P1,
                signals=[Signal(
                    alert_name="TestAlert",
                    severity=Severity.P1,
                    source="test",
                    fired_at=datetime.now(timezone.utc),
                )],
                affected_services=["test-service"],
            )

        async def update_incident(self, incident_id, updates):
            return True

    engine = ProposalEngine(memory=MockMemory())

    # 取得預設 Guardrails
    guardrails = engine.get_default_guardrails()

    assert guardrails.require_dry_run == True
    assert "awoooi-prod" in guardrails.allowed_namespace
    assert any("rm -rf" in cmd.lower() for cmd in guardrails.forbidden_commands)

    print("✅ ProposalEngine Guardrails 整合成功:")
    print(f"   - require_dry_run: {guardrails.require_dry_run}")
    print(f"   - allowed_namespace: {guardrails.allowed_namespace}")
    print(f"   - forbidden_commands: {len(guardrails.forbidden_commands)} 項")


if __name__ == "__main__":
    print("=" * 60)
    print("🧪 Guardrails 單元測試")
    print("=" * 60)

    tests = [
        test_guardrails_validator_import,
        test_guardrails_block_rm_rf,
        test_guardrails_block_drop_database,
        test_guardrails_block_kubectl_delete_ns,
        test_guardrails_allow_safe_commands,
        test_guardrails_namespace_validation,
        test_guardrails_enforce_dry_run,
        test_proposal_engine_guardrails_integration,
    ]

    passed = 0
    failed = 0

    for test in tests:
        print(f"\n🔬 {test.__name__}")
        try:
            test()
            passed += 1
        except AssertionError as e:
            print(f"❌ FAILED: {e}")
            failed += 1
        except Exception as e:
            print(f"❌ ERROR: {type(e).__name__}: {e}")
            failed += 1

    print("\n" + "=" * 60)
    print(f"📊 結果: {passed} 通過, {failed} 失敗")
    print("=" * 60)

    if failed > 0:
        sys.exit(1)