HIPAA Compliance Validation: Audit Trail and Key Management

import hashlib
import json
import uuid
from datetime import datetime, timezone
from dataclasses import dataclass, field
from typing import List, Optional
from mirror_sdk.core.mirror_core import MirrorSDK, MirrorConfig
from mirror_sdk.core.mirror_api_models import Action

config = MirrorConfig.from_env()
sdk = MirrorSDK(config)

@dataclass
class AuditEvent:
    user_id: str
    user_role: str
    department: str
    action: str
    record_ids_accessed: List[str]
    query_hash: str
    session_id: str
    response_phi_detected: bool = False
    event_id: str = field(default_factory=lambda: f"evt_{uuid.uuid4().hex[:10]}")
    timestamp_utc: str = field(
        default_factory=lambda: datetime.now(timezone.utc).isoformat()
    )
    log_hash: str = ""

    def compute_hash(self) -> str:
        """Compute a tamper-detection hash of this log entry."""
        payload = {
            "event_id": self.event_id,
            "timestamp_utc": self.timestamp_utc,
            "user_id": self.user_id,
            "user_role": self.user_role,
            "department": self.department,
            "action": self.action,
            "record_ids_accessed": self.record_ids_accessed,
            "query_hash": self.query_hash,
            "session_id": self.session_id,
            "response_phi_detected": self.response_phi_detected,
        }
        canonical = json.dumps(payload, sort_keys=True)
        return hashlib.sha256(canonical.encode()).hexdigest()

def log_phi_access(
    sdk: MirrorSDK,
    user_id: str,
    user_role: str,
    department: str,
    query_text: str,
    record_ids: List[str],
    model_response: str,
    session_id: str
) -> AuditEvent:
    """
    Log a PHI access event with tamper detection.
    Checks the model response for PHI leakage before writing.
    Does not store the query text or response text in the log.
    """

    # Hash the query (do not store the query itself)
    query_hash = "sha256:" + hashlib.sha256(query_text.encode()).hexdigest()

    # Check model response for PHI leakage using AgentIQ
    phi_result = sdk.agentiq.detect_pii(
        text=model_response,
        pii_entities=["NAME", "EMAIL", "PHONE", "SSN", "MRN", "DATE"],
        action=Action.ALERT
    )
    phi_detected = len(phi_result.entities) > 0

    # Build the audit event
    event = AuditEvent(
        user_id=user_id,
        user_role=user_role,
        department=department,
        action="rag_query",
        record_ids_accessed=record_ids,
        query_hash=query_hash,
        session_id=session_id,
        response_phi_detected=phi_detected
    )

    # Compute tamper-detection hash and write it into the event
    event.log_hash = event.compute_hash()

    # Write to your audit log store
    # This should be an append-only, write-once store (e.g. Azure Immutable Blob Storage)
    write_audit_log(event)

    if phi_detected:
        trigger_phi_leakage_alert(event, phi_result.entities)

    return event

import json
import hashlib

def verify_audit_log_integrity(log_entries: list) -> dict:
    """
    Verify that no audit log entries have been tampered with.
    Returns a report of any integrity failures.
    """
    failures = []

    for entry in log_entries:
        stored_hash = entry.get("log_hash", "")

        # Recompute the hash from the entry fields
        check_fields = {k: v for k, v in entry.items() if k != "log_hash"}
        canonical = json.dumps(check_fields, sort_keys=True)
        expected_hash = hashlib.sha256(canonical.encode()).hexdigest()

        if stored_hash != expected_hash:
            failures.append({
                "event_id": entry.get("event_id"),
                "timestamp_utc": entry.get("timestamp_utc"),
                "issue": "hash_mismatch"
            })

    return {
        "total_entries": len(log_entries),
        "failures": failures,
        "integrity_verified": len(failures) == 0
    }

from mirror_sdk.core.mirror_core import MirrorSDK, MirrorConfig
import logging

logger = logging.getLogger("hipaa_key_rotation")
config = MirrorConfig.from_env()
sdk = MirrorSDK(config)

def rotate_fpe_key(
    metadata_records: list,
    old_key: str,
    key_vault_client
) -> dict:
    """
    Rotate the FPE key for patient metadata.
    1. Generate a new key in the key vault
    2. Re-encrypt all metadata records with the new key
    3. Archive the old key (retain for 6 years for historical decryption)
    4. Log the rotation event for compliance
    """

    # Step 1: Generate new key
    new_key = sdk.metadata.generate_key()
    new_key_id = key_vault_client.store_key(
        key=new_key,
        label="vectax-fpe-active"
    )
    logger.info(f"New FPE key generated: {new_key_id}")

    # Step 2: Re-encrypt all metadata records
    reencrypted = []
    failed = []

    for record in metadata_records:
        try:
            # Decrypt with old key
            tweak = sdk.metadata.generate_tweak_from_data(record["encrypted_metadata"])
            plaintext_metadata = sdk.metadata.decrypt(
                record["encrypted_metadata"], old_key, tweak
            )

            # Re-encrypt with new key
            new_tweak = sdk.metadata.generate_tweak_from_data(plaintext_metadata)
            new_encrypted = sdk.metadata.encrypt(plaintext_metadata, new_key, new_tweak)

            reencrypted.append({
                "record_id": record["record_id"],
                "encrypted_metadata": new_encrypted,
                "key_id": new_key_id
            })

        except Exception as e:
            logger.error(f"Failed to re-encrypt record {record['record_id']}: {e}")
            failed.append(record["record_id"])

    # Step 3: Archive (do not delete) the old key
    key_vault_client.archive_key(old_key, label="vectax-fpe-retired")

    # Step 4: Log the rotation for compliance evidence
    logger.info(
        f"FPE key rotation complete. "
        f"Records re-encrypted: {len(reencrypted)}. "
        f"Failures: {len(failed)}. "
        f"New key ID: {new_key_id}"
    )

    return {
        "new_key_id": new_key_id,
        "reencrypted_count": len(reencrypted),
        "failed_records": failed
    }

from mirror_sdk.core.mirror_core import MirrorSDK, MirrorConfig
from mirror_sdk.core.mirror_api_models import Action
import logging

logger = logging.getLogger("phi_output_guard")
config = MirrorConfig.from_env()
sdk = MirrorSDK(config)

# PHI entity types relevant to clinical AI outputs
CLINICAL_PHI_ENTITIES = [
    "NAME",         # Patient name
    "DATE",         # Dates of service, birth dates
    "PHONE",        # Phone numbers
    "EMAIL",        # Email addresses
    "SSN",          # Social security numbers
    "MRN",          # Medical record numbers
    "ADDRESS",      # Geographic identifiers
    "MEDICAL_RECORD" # Diagnosis codes, procedure codes
]

def guard_clinical_output(
    model_response: str,
    action: Action = Action.REDACT
) -> dict:
    """
    Scan a clinical AI model response for PHI.
    REDACT by default: returns the response with PHI replaced.
    ALERT: returns the response unchanged but flags for review.
    BLOCK: raises an exception if PHI is detected.
    """

    result = sdk.agentiq.detect_pii(
        text=model_response,
        pii_entities=CLINICAL_PHI_ENTITIES,
        action=action
    )

    if result.entities:
        logger.warning(
            f"PHI detected in clinical AI output. "
            f"Entities: {[e.label for e in result.entities]}. "
            f"Action: {action}. "
            f"Risk score: {result.risk_score}"
        )

    return {
        "original_length": len(model_response),
        "phi_detected": len(result.entities) > 0,
        "entity_count": len(result.entities),
        "entity_types": [e.label for e in result.entities],
        "risk_score": result.risk_score,
        # With REDACT action, this is the safe version to return to the user
        "safe_response": result.redacted_text if action == Action.REDACT else model_response
    }