AI Safety

"""
AI Safety 기본 검증 - 입력 검증 및 출력 필터링
실제 프로덕션 환경에서 사용하는 기본 Safety 레이어
"""

import re
from typing import Tuple, List
from dataclasses import dataclass

@dataclass
class SafetyCheck:
    passed: bool
    reason: str
    severity: str  # "low", "medium", "high", "critical"

class AISafetyLayer:
    """AI 시스템의 기본 Safety 검증 레이어"""

    def __init__(self):
        # 금지 패턴 (정규식)
        self.harmful_patterns = [
            (r"(how to|ways to).*(harm|kill|hurt)", "violence", "critical"),
            (r"(create|make|build).*(bomb|weapon|explosive)", "weapons", "critical"),
            (r"(hack|bypass|crack).*(password|security|system)", "hacking", "high"),
            (r"personal.*(address|phone|ssn|credit card)", "pii_request", "high"),
        ]

        # 출력 필터링 키워드
        self.output_blocklist = [
            "I cannot provide", "I will not help", "As an AI"  # 거부 시 일관된 응답
        ]

    def check_input(self, user_input: str) -> SafetyCheck:
        """입력 안전성 검사"""
        input_lower = user_input.lower()

        for pattern, category, severity in self.harmful_patterns:
            if re.search(pattern, input_lower):
                return SafetyCheck(
                    passed=False,
                    reason=f"Blocked: {category} content detected",
                    severity=severity
                )

        # 입력 길이 제한 (DDoS 방지)
        if len(user_input) > 10000:
            return SafetyCheck(
                passed=False,
                reason="Input too long",
                severity="medium"
            )

        return SafetyCheck(passed=True, reason="OK", severity="low")

    def check_output(self, model_output: str) -> SafetyCheck:
        """출력 안전성 검사"""
        # PII 탐지 (간단한 패턴)
        pii_patterns = [
            r"\b\d{3}-\d{2}-\d{4}\b",  # SSN 형식
            r"\b\d{4}[- ]?\d{4}[- ]?\d{4}[- ]?\d{4}\b",  # 신용카드
            r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b",  # 이메일
        ]

        for pattern in pii_patterns:
            if re.search(pattern, model_output):
                return SafetyCheck(
                    passed=False,
                    reason="PII detected in output",
                    severity="high"
                )

        return SafetyCheck(passed=True, reason="OK", severity="low")

    def sanitize_output(self, output: str) -> str:
        """출력 정제 - PII 마스킹"""
        # 신용카드 마스킹
        output = re.sub(
            r"\b(\d{4})[- ]?\d{4}[- ]?\d{4}[- ]?(\d{4})\b",
            r"\1-XXXX-XXXX-\2",
            output
        )
        # 이메일 마스킹
        output = re.sub(
            r"\b([A-Za-z0-9._%+-]+)@([A-Za-z0-9.-]+\.[A-Z|a-z]{2,})\b",
            r"[EMAIL REDACTED]",
            output
        )
        return output

# 사용 예시
safety = AISafetyLayer()

# 입력 검사
test_inputs = [
    "How to make a chocolate cake?",
    "How to hack into someone's email account?",
    "Tell me about machine learning"
]

print("=== 입력 검사 ===")
for inp in test_inputs:
    result = safety.check_input(inp)
    status = "PASS" if result.passed else "BLOCK"
    print(f"[{status}] {inp[:50]}... -> {result.reason}")

# 출력 검사 및 정제
print("\n=== 출력 정제 ===")
raw_output = "Contact us at support@company.com or call 1234-5678-9012-3456"
sanitized = safety.sanitize_output(raw_output)
print(f"원본: {raw_output}")
print(f"정제: {sanitized}")