1
from dataclasses import dataclass, field, asdict
2
from datetime import datetime
3
from typing import Any
4
import json
5
import uuid
6

7
@dataclass
8
class AgentAuditRecord:
9
    """Complete audit record for a single agent interaction."""
10

11
    # Identity
12
    audit_id: str = field(default_factory=lambda: uuid.uuid4().hex)
13
    agent_id: str
14
    agent_version: str           # Prompt version, steering file hash
15
    session_id: str
16
    user_id: str
17

18
    # Input
19
    timestamp: datetime = field(default_factory=datetime.utcnow)
20
    raw_input: str               # Original user message
21
    sanitized_input: str         # PII-redacted version
22
    input_guardrail_results: list
23

24
    # Decision trace (for explainability)
25
    system_prompt_used: str
26
    steering_files_loaded: list[str]
27
    mcp_servers_connected: list[str]
28

29
    # Tool calls
30
    tool_calls: list[dict] = field(default_factory=list)
31
    # Each tool call record:
32
    # {
33
    #   "tool": "query_database",
34
    #   "params": {"sql": "SELECT..."},
35
    #   "params_redacted": {"sql": "SELECT..."},  # PII removed
36
    #   "result_preview": "3 rows returned",
37
    #   "result_redacted": [...],
38
    #   "duration_ms": 234,
39
    #   "success": True,
40
    #   "guardrail_check": "passed",
41
    # }
42

43
    # Output
44
    raw_output: str
45
    sanitized_output: str
46
    output_guardrail_results: list
47

48
    # Compliance metadata
49
    data_retention_days: int = 90
50
    contains_pii: bool
51
    contains_pci: bool
52
    retention_policy: str       # "standard", "pii_retention", "pci_retention"
53

54
    # Human-in-loop
55
    human_approvals: list[dict] = field(default_factory=list)
56
    # {
57
    #   "tool": "deploy_production",
58
    #   "approved_by": "user_abc",
59
    #   "approved_at": "2026-06-04T12:00:00Z",
60
    #   "reasoning_summary": "Agent explained X, Y, Z"
61
    # }
62

63
    def to_storage_format(self) -> dict:
64
        """Convert to format suitable for encrypted storage."""
65
        record = asdict(self)
66
        record["timestamp"] = self.timestamp.isoformat()
67
        return record
68

69
class AuditTrailLogger:
70
    """Secure audit trail logger with encryption and immutability."""
71

72
    def __init__(self, storage_backend="s3", encryption_key=None):
73
        self.storage = AuditLogStorage(storage_backend)
74
        self.encryptor = FieldEncryptor(encryption_key) if encryption_key else None
75

76
    async def log_interaction(self, record: AgentAuditRecord):
77
        # Encrypt sensitive fields
78
        if self.encryptor:
79
            record.raw_input = self.encryptor.encrypt(record.raw_input)
80
            record.raw_output = self.encryptor.encrypt(record.raw_output)
81

82
        # Add hash chain (immutability)
83
        prev_hash = await self.storage.get_last_hash()
84
        record.prev_hash = prev_hash
85
        record.hash = self.compute_hash(record)
86

87
        # Store
88
        await self.storage.append(f"agent/{record.agent_id}/{record.session_id}", record)
89

90
    def compute_hash(self, record: AgentAuditRecord) -> str:
91
        content = json.dumps(record.to_storage_format(), sort_keys=True)
92
        if record.prev_hash:
93
            content = record.prev_hash + content
94
        return hashlib.sha256(content.encode()).hexdigest()
95

96
    async def get_session_trail(self, session_id: str) -> list[AgentAuditRecord]:
97
        """Retrieve complete audit trail for a session."""
98
        return await self.storage.query(session_id=session_id)

1
class ReasoningCapturer:
2
    """Capture the agent's reasoning process for explainability."""
3

4
    def __init__(self):
5
        self.reasoning_log = []
6

7
    async def capture_trace(self, agent_run):
8
        """Wrap an agent run to capture reasoning."""
9

10
        # Hook into agent events
11
        agent_run.on("thinking", self.capture_thinking)
12
        agent_run.on("tool_call", self.capture_tool_decision)
13
        agent_run.on("tool_result", self.capture_tool_result)
14
        agent_run.on("final_output", self.capture_final_output)
15

16
    def capture_thinking(self, thought: str):
17
        self.reasoning_log.append({
18
            "type": "thought",
19
            "timestamp": datetime.utcnow().isoformat(),
20
            "content": thought,
21
        })
22

23
    def capture_tool_decision(self, tool: str, args: dict):
24
        self.reasoning_log.append({
25
            "type": "tool_decision",
26
            "timestamp": datetime.utcnow().isoformat(),
27
            "tool": tool,
28
            "args": self.redact_pii(args),
29
            "reasoning_context": self.get_current_context(),
30
        })
31

32
    def capture_tool_result(self, tool: str, result: Any):
33
        self.reasoning_log.append({
34
            "type": "tool_result",
35
            "timestamp": datetime.utcnow().isoformat(),
36
            "tool": tool,
37
            "result_preview": str(result)[:200] + "...",
38
        })
39

40
    def get_explainability_report(self) -> dict:
41
        """Generate a human-readable explanation of the agent's decisions."""
42
        return {
43
            "chain_of_thought": [
44
                entry for entry in self.reasoning_log
45
                if entry["type"] == "thought"
46
            ],
47
            "tool_calls": [
48
                {
49
                    "step": i,
50
                    "tool": entry["tool"],
51
                    "why": self.infer_tool_reasoning(entry),
52
                }
53
                for i, entry in enumerate(self.reasoning_log)
54
                if entry["type"] == "tool_decision"
55
            ],
56
            "summary": self.generate_summary(),
57
        }
58

59
    def infer_tool_reasoning(self, entry: dict) -> str:
60
        """From the surrounding thoughts, infer why this tool was called."""
61
        # Find the thoughts before this tool call
62
        idx = self.reasoning_log.index(entry)
63
        preceding_thoughts = [
64
            e["content"] for e in self.reasoning_log[max(0, idx-3):idx]
65
            if e["type"] == "thought"
66
        ]
67
        return " — ".join(preceding_thoughts) if preceding_thoughts else "No reasoning captured"

1
class GDPRCompliance:
2
    """Handle GDPR right to explanation requests."""
3

4
    async def generate_explanation(
5
        self,
6
        user_id: str,
7
        session_id: str,
8
        audit_trail: AuditTrailLogger,
9
    ) -> dict:
10
        """Generate a human-readable explanation for GDPR Article 22 compliance."""
11

12
        records = await audit_trail.get_session_trail(session_id)
13

14
        explanation = {
15
            "request_date": datetime.utcnow().isoformat(),
16
            "data_subject": user_id,
17
            "automated_decision": {
18
                "made": True,
19
                "logic_involved": "Large Language Model with tool access",
20
                "significance": self.describe_significance(records),
21
            },
22
            "decision_sequence": [],
23
        }
24

25
        for record in records:
26
            step = {
27
                "input": record.sanitized_input[:500],
28
                "tools_used": [tc["tool"] for tc in record.tool_calls],
29
                "reasoning": self.extract_reasoning(record),
30
                "output_preview": record.sanitized_output[:500],
31
            }
32
            explanation["decision_sequence"].append(step)
33

34
        return explanation
35

36
    async def handle_deletion_request(self, user_id: str):
37
        """GDPR right to erasure."""
38
        await self.anonymize_user_data(user_id)
39
        await self.audit_log.log_deletion(user_id)

1
import re
2
from presidio_analyzer import AnalyzerEngine
3
from presidio_anonymizer import AnonymizerEngine
4

5
class PIIHandler:
6
    """Detect, redact, and manage PII in agent interactions."""
7

8
    def __init__(self):
9
        self.analyzer = AnalyzerEngine()
10
        self.anonymizer = AnonymizerEngine()
11

12
        # Regex patterns for fast pre-filtering
13
        self.patterns = {
14
            "email": r'\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b',
15
            "phone": r'\b\+?\d{1,3}[-.\s]?\(?\d{1,4}\)?[-.\s]?\d{1,4}[-.\s]?\d{1,9}\b',
16
            "ssn": r'\b\d{3}-\d{2}-\d{4}\b',
17
            "credit_card": r'\b(?:\d{4}[-\s]?){3}\d{4}\b',
18
        }
19

20
    async def process_input(self, text: str, context: str = "") -> ProcessedContent:
21
        """Analyze and redact PII from agent input."""
22

23
        # 1. Fast regex detection
24
        found_patterns = {}
25
        for pattern_name, pattern in self.patterns.items():
26
            matches = re.findall(pattern, text)
27
            if matches:
28
                found_patterns[pattern_name] = matches
29

30
        # 2. NLP-based detection (catches context-sensitive PII)
31
        analyzer_results = self.analyzer.analyze(text=text, language='en')
32

33
        # 3. Redact
34
        redacted = self.anonymizer.anonymize(
35
            text=text,
36
            analyzer_results=analyzer_results,
37
        )
38

39
        return ProcessedContent(
40
            original=text,
41
            redacted=redacted.text,
42
            detected_pii=found_patterns,
43
            has_pii=len(found_patterns) > 0 or len(analyzer_results) > 0,
44
            pii_categories=list(set(
45
                list(found_patterns.keys()) +
46
                [r.entity_type for r in analyzer_results]
47
            )),
48
        )
49

50
    async def handle_pci_data(self, text: str) -> PCIResult:
51
        """Special handling for PCI DSS cardholder data."""
52

53
        credit_card_pattern = r'\b(?:\d{4}[-\s]?){3}\d{4}\b'
54
        if re.search(credit_card_pattern, text):
55
            # PCI DSS Requirement 3.4: Render PAN unreadable
56
            redacted = re.sub(credit_card_pattern, '****-****-****-****', text)
57

58
            # Log PCI event (separate from standard audit)
59
            await self.pci_audit_log.log_pci_event(
60
                event_type="card_data_detected",
61
                action="redacted",
62
                timestamp=datetime.utcnow(),
63
            )
64

65
            # Never store full PAN
66
            return PCIResult(
67
                original_redacted=True,
68
                safe_output=redacted,
69
                pci_audit_id=uuid.uuid4().hex,
70
            )
71

72
        return PCIResult(original_redacted=False, safe_output=text)

1
class DataRetentionManager:
2
    """Manage data retention based on data classification."""
3

4
    RETENTION_RULES = {
5
        "standard": {
6
            "conversation_logs": timedelta(days=90),
7
            "tool_call_logs": timedelta(days=90),
8
            "aggregate_metrics": timedelta(days=365),
9
        },
10
        "pii": {
11
            "conversation_logs": timedelta(days=30),  # Shorter for PII
12
            "tool_call_logs": timedelta(days=90),
13
            "pii_detection_log": timedelta(days=365),
14
        },
15
        "pci": {
16
            "conversation_logs": timedelta(days=0),   # Never store raw
17
            "tool_call_logs": timedelta(days=365),      # Longer for audit
18
            "pci_event_log": timedelta(days=2555),      # 7 years for PCI
19
        },
20
    }
21

22
    async def apply_retention_policy(self):
23
        """Apply retention policies to all stored data."""
24
        for data_class, rules in self.RETENTION_RULES.items():
25
            for log_type, retention in rules.items():
26
                cutoff = datetime.utcnow() - retention
27
                await self.storage.delete_older_than(
28
                    collection=f"{data_class}/{log_type}",
29
                    cutoff=cutoff,
30
                )
31

32
    async def handle_deletion_request(self, user_id: str):
33
        """GDPR Article 17: Right to erasure."""
34
        # Anonymize user-specific audit records
35
        await self.storage.anonymize_field(
36
            collection="standard/conversation_logs",
37
            field="user_id",
38
            value=user_id,
39
            replacement="ANONYMIZED",
40
        )
41
        # Delete PII-specific records
42
        await self.storage.delete(
43
            collection="pii/conversation_logs",
44
            where={"user_id": user_id},
45
        )

1
class SOC2EvidenceCollector:
2
    """Generate evidence for SOC 2 audits."""
3

4
    async def collect_controls_evidence(self, start_date, end_date):
5
        return {
6
            "cc6_1_access_control": {
7
                "control": "MCP servers require authentication",
8
                "evidence": await self.get_auth_logs(start_date, end_date),
9
                "pass_rate": await self.calc_auth_pass_rate(),
10
                "exceptions": await self.get_auth_failures(),
11
            },
12
            "cc6_6_injection_prevention": {
13
                "control": "Input guardrails block injection attempts",
14
                "evidence": await self.get_guardrail_logs(),
15
                "blocked_count": await self.count_blocked_injections(),
16
                "false_positive_rate": await self.calc_false_positive_rate(),
17
            },
18
            "cc7_2_monitoring": {
19
                "control": "All agent interactions are logged",
20
                "evidence": await self.get_audit_log_coverage(),
21
                "coverage_percentage": await self.calc_log_coverage(),
22
            },
23
            "a1_2_processing_integrity": {
24
                "control": "Tool calls are validated before execution",
25
                "evidence": await self.get_tool_validation_logs(),
26
                "validation_pass_rate": await self.calc_validation_rate(),
27
            },
28
        }