
# compliance/mapper.py — Gap #3: Regulatory Standards Mapper.

# Automatically maps vulnerability types to regulatory standards:
#   EU AI Act (Articles 9, 10, 13, 15, Annex III)
#    NIST AI RMF (GOVERN, MAP, MEASURE, MANAGE functions)
#  HIPAA (45 CFR §164.306, §164.312)
#  GDPR (Articles 5, 22, 25, 32)
#  OWASP LLM Top 10 (LLM01–LLM10)

# 
#     mapper = ComplianceMapper()
#     mapping = mapper.map_finding(attack_result_dict)
from __future__ import annotations

from typing import Dict, List, Optional

#standar ref
# EU AI Act — Article citations relevant to LLM security testing
_EU_AI_ACT: Dict[str, List[str]] = {
    "jailbreak":          ["Article 9.2", "Article 15.1", "Annex III"],
    "prompt_injection":   ["Article 9.2", "Article 15.1", "Article 15.3"],
    "data_leakage":       ["Article 10.3", "Article 15.2", "Article 13.1"],
    "data_extraction":    ["Article 10.3", "Article 15.2"],
    "hallucination":      ["Article 13.1", "Article 13.2", "Article 15.4"],
    "bias_fairness":      ["Article 9.7", "Article 10.2", "Annex III"],
    "alignment":          ["Article 9.2", "Article 15.1"],
    "backdoor":           ["Article 9.4", "Article 15.1", "Annex III"],
    "system_prompt_leak": ["Article 13.1", "Article 15.3"],
    "unknown":            ["Article 9.2"],
}

# NIST AI RMF — Function + Sub-category references
_NIST_AI_RMF: Dict[str, List[str]] = {
    "jailbreak":          ["GOVERN-1.1", "GOVERN-1.7", "MEASURE-2.3", "MANAGE-2.2"],
    "prompt_injection":   ["GOVERN-1.1", "MAP-2.3", "MEASURE-2.3", "MANAGE-2.2"],
    "data_leakage":       ["GOVERN-1.1", "MAP-1.6", "MEASURE-2.5", "MANAGE-3.1"],
    "data_extraction":    ["MAP-1.6", "MEASURE-2.5", "MANAGE-3.1"],
    "hallucination":      ["GOVERN-4.1", "MAP-5.1", "MEASURE-2.6"],
    "bias_fairness":      ["GOVERN-4.1", "MAP-5.2", "MEASURE-2.2", "MANAGE-2.4"],
    "alignment":          ["GOVERN-1.7", "MEASURE-2.3"],
    "backdoor":           ["GOVERN-1.1", "MAP-2.3", "MEASURE-2.3"],
    "system_prompt_leak": ["GOVERN-1.7", "MAP-2.3"],
    "unknown":            ["GOVERN-1.1", "MEASURE-2.3"],
}

# HIPAA — 45 CFR citations
_HIPAA: Dict[str, List[str]] = {
    "jailbreak":          ["45 CFR 164.306(a)", "45 CFR 164.312(a)(1)"],
    "prompt_injection":   ["45 CFR 164.306(a)", "45 CFR 164.312(a)(1)"],
    "data_leakage":       ["45 CFR 164.312(b)", "45 CFR 164.314(a)(2)"],
    "data_extraction":    ["45 CFR 164.312(b)", "45 CFR 164.314(a)(2)"],
    "hallucination":      ["45 CFR 164.306(a)"],
    "bias_fairness":      ["45 CFR 164.306(a)"],
    "alignment":          ["45 CFR 164.306(a)"],
    "backdoor":           ["45 CFR 164.306(a)", "45 CFR 164.312(a)(1)"],
    "system_prompt_leak": ["45 CFR 164.312(b)"],
    "unknown":            [],
}

# GDPR — Article citations
_GDPR: Dict[str, List[str]] = {
    "jailbreak":          ["Article 5", "Article 32"],
    "prompt_injection":   ["Article 5", "Article 32"],
    "data_leakage":       ["Article 5", "Article 25", "Article 32", "Article 33"],
    "data_extraction":    ["Article 5", "Article 25", "Article 32"],
    "hallucination":      ["Article 22"],
    "bias_fairness":      ["Article 22", "Article 5(1)(f)"],
    "alignment":          ["Article 22"],
    "backdoor":           ["Article 5", "Article 32"],
    "system_prompt_leak": ["Article 5", "Article 25"],
    "unknown":            ["Article 32"],
}

# OWASP LLM Top 10 (2025 edition)
_OWASP_LLM: Dict[str, List[str]] = {
    "jailbreak":          ["LLM01:2025 Prompt Injection"],
    "prompt_injection":   ["LLM01:2025 Prompt Injection"],
    "data_leakage":       ["LLM06:2025 Sensitive Information Disclosure", "LLM02:2025 Sensitive Information Disclosure"],
    "data_extraction":    ["LLM10:2025 Unbounded Consumption", "LLM06:2025 Sensitive Information Disclosure"],
    "hallucination":      ["LLM09:2025 Misinformation"],
    "bias_fairness":      ["LLM09:2025 Misinformation"],
    "alignment":          ["LLM01:2025 Prompt Injection"],
    "backdoor":           ["LLM05:2025 Improper Output Handling"],
    "system_prompt_leak": ["LLM07:2025 System Prompt Leakage"],
    "unknown":            [],
}


#compliance mapper

class ComplianceMapper:
    

    def map_finding(self, result: dict) -> Dict[str, List[str]]:
        
        from sichgate.compliance.models import _infer_vulnerability_type
        vuln_type = _infer_vulnerability_type(result)
        return self.map_vulnerability_type(vuln_type)

    def map_vulnerability_type(self, vulnerability_type: str) -> Dict[str, List[str]]:
        
        vt = vulnerability_type if vulnerability_type in _EU_AI_ACT else "unknown"
        return {
            "eu_ai_act":    _EU_AI_ACT.get(vt, _EU_AI_ACT["unknown"]),
            "nist_ai_rmf":  _NIST_AI_RMF.get(vt, _NIST_AI_RMF["unknown"]),
            "hipaa":        _HIPAA.get(vt, []),
            "gdpr":         _GDPR.get(vt, _GDPR["unknown"]),
            "owasp":        _OWASP_LLM.get(vt, []),
        }

    def map_run(self, run_summary: dict) -> List[dict]:
        
        from sichgate.compliance.models import ComplianceAuditTrail
        run_id = run_summary.get("run_id")
        findings = []

        for result in run_summary.get("results", []):
            # Only create audit entries for failed/warning results
            if result.get("status") not in ("fail", "warning"):
                continue

            mapping = self.map_finding(result)
            trail = ComplianceAuditTrail.from_attack_result(
                result, run_id=run_id, standards_mapping=mapping
            )
            findings.append({
                "trail": trail.to_dict(),
                "standards": mapping,
            })

        return findings

    def list_supported_vulnerability_types(self) -> List[str]:
        
        return list(_EU_AI_ACT.keys())

    def get_standards_reference(self, vulnerability_type: str) -> Dict[str, dict]:
        
        mapping = self.map_vulnerability_type(vulnerability_type)
        return {
            "EU AI Act": {
                "articles": mapping["eu_ai_act"],
                "url": "https://eur-lex.europa.eu/legal-content/EN/TXT/?uri=CELEX:32024R1689",
            },
            "NIST AI RMF": {
                "functions": mapping["nist_ai_rmf"],
                "url": "https://airc.nist.gov/RMF/RMF_Documentation",
            },
            "HIPAA": {
                "rules": mapping["hipaa"],
                "url": "https://www.hhs.gov/hipaa/for-professionals/security/index.html",
            },
            "GDPR": {
                "articles": mapping["gdpr"],
                "url": "https://gdpr-info.eu/",
            },
            "OWASP LLM Top 10": {
                "references": mapping["owasp"],
                "url": "https://owasp.org/www-project-top-10-for-large-language-model-applications/",
            },
        }
