zoey_storage_sql/
hipaa.rs

1//! HIPAA Compliance Features for Database
2//!
3//! Provides database-level protection for PII including:
4//! - Pre-storage PII detection and redaction
5//! - Dual storage (encrypted original + redacted version)
6//! - Access control with attorney-client privilege support
7//! - Comprehensive audit logging
8//! - HIPAA-compliant encryption at rest
9
10use zoey_core::{
11    security::{decrypt_secret, encrypt_secret},
12    ZoeyError, Result,
13};
14use serde::{Deserialize, Serialize};
15use sqlx::PgPool;
16use std::collections::HashMap;
17
18/// HIPAA compliance configuration
19#[derive(Debug, Clone)]
20pub struct HIPAAConfig {
21    /// Enable HIPAA features (master switch)
22    pub enabled: bool,
23
24    /// Enable audit logging
25    pub audit_logging: bool,
26
27    /// Enable encryption at rest
28    pub encryption_at_rest: bool,
29
30    /// Enable access control
31    pub access_control: bool,
32
33    /// Retention policy in days
34    pub retention_days: usize,
35
36    /// Enable automatic de-identification
37    pub auto_deidentify: bool,
38
39    /// Encryption key for PHI data (MUST be set for encryption_at_rest)
40    /// This should be loaded from a secure key management system, NOT hardcoded
41    encryption_key: Option<String>,
42}
43
44impl Default for HIPAAConfig {
45    fn default() -> Self {
46        Self {
47            enabled: true, // Can be disabled
48            audit_logging: true,
49            encryption_at_rest: true,
50            access_control: true,
51            retention_days: 2555, // 7 years as per HIPAA
52            auto_deidentify: true,
53            encryption_key: None, // MUST be set via with_encryption_key() for production
54        }
55    }
56}
57
58impl HIPAAConfig {
59    /// Create config with HIPAA disabled (for non-healthcare use)
60    pub fn disabled() -> Self {
61        Self {
62            enabled: false,
63            audit_logging: false,
64            encryption_at_rest: false,
65            access_control: false,
66            retention_days: 365, // 1 year default
67            auto_deidentify: false,
68            encryption_key: None,
69        }
70    }
71
72    /// Create config with minimal features (audit only)
73    pub fn minimal() -> Self {
74        Self {
75            enabled: true,
76            audit_logging: true, // Keep audit trail
77            encryption_at_rest: false,
78            access_control: false,
79            retention_days: 365,
80            auto_deidentify: false,
81            encryption_key: None,
82        }
83    }
84
85    /// Create config with maximum compliance
86    pub fn maximum() -> Self {
87        Self::default()
88    }
89
90    /// Set the encryption key for PHI data
91    ///
92    /// SECURITY: The key should be:
93    /// - At least 32 characters long
94    /// - Loaded from a secure key management system (e.g., AWS KMS, HashiCorp Vault)
95    /// - NEVER hardcoded in source code
96    /// - Rotated periodically
97    pub fn with_encryption_key(mut self, key: impl Into<String>) -> Self {
98        let key = key.into();
99        if key.len() < 32 {
100            tracing::warn!(
101                "HIPAA encryption key is less than 32 characters - this may be insecure"
102            );
103        }
104        self.encryption_key = Some(key);
105        self
106    }
107
108    /// Check if encryption is properly configured
109    pub fn is_encryption_ready(&self) -> bool {
110        self.encryption_at_rest && self.encryption_key.is_some()
111    }
112}
113
114/// HIPAA audit log entry
115#[derive(Debug, Clone, Serialize, Deserialize)]
116pub struct AuditLogEntry {
117    /// Unique identifier
118    pub id: uuid::Uuid,
119
120    /// Timestamp
121    pub timestamp: i64,
122
123    /// User/entity who performed the action
124    pub actor_id: uuid::Uuid,
125
126    /// Action performed
127    pub action: String,
128
129    /// Resource accessed
130    pub resource_type: String,
131
132    /// Resource ID
133    pub resource_id: uuid::Uuid,
134
135    /// IP address
136    pub ip_address: Option<String>,
137
138    /// Result (success/failure)
139    pub result: String,
140
141    /// Additional metadata
142    pub metadata: HashMap<String, String>,
143}
144
145/// HIPAA compliance manager
146pub struct HIPAACompliance {
147    pool: PgPool,
148    config: HIPAAConfig,
149}
150
151impl HIPAACompliance {
152    /// Create a new HIPAA compliance manager
153    pub fn new(pool: PgPool, config: HIPAAConfig) -> Self {
154        Self { pool, config }
155    }
156
157    /// Initialize HIPAA compliance tables
158    pub async fn initialize(&self) -> Result<()> {
159        if !self.config.enabled {
160            tracing::info!("HIPAA compliance is DISABLED - skipping initialization");
161            return Ok(());
162        }
163
164        tracing::info!("Initializing HIPAA compliance features");
165
166        // Create audit log table
167        if self.config.audit_logging {
168            sqlx::query(
169                r#"
170            CREATE TABLE IF NOT EXISTS audit_log (
171                id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
172                timestamp BIGINT NOT NULL,
173                actor_id UUID NOT NULL,
174                action TEXT NOT NULL,
175                resource_type TEXT NOT NULL,
176                resource_id UUID NOT NULL,
177                ip_address TEXT,
178                result TEXT NOT NULL,
179                metadata JSONB DEFAULT '{}'
180            )
181        "#,
182            )
183            .execute(&self.pool)
184            .await?;
185
186            // Create index for audit log queries
187            sqlx::query(
188                "CREATE INDEX IF NOT EXISTS idx_audit_log_timestamp ON audit_log(timestamp)",
189            )
190            .execute(&self.pool)
191            .await?;
192
193            sqlx::query("CREATE INDEX IF NOT EXISTS idx_audit_log_actor ON audit_log(actor_id)")
194                .execute(&self.pool)
195                .await?;
196        }
197
198        // Enable Row Level Security (RLS)
199        if self.config.access_control {
200            sqlx::query("ALTER TABLE memories ENABLE ROW LEVEL SECURITY")
201                .execute(&self.pool)
202                .await
203                .ok(); // May fail if already enabled
204
205            // Create policies for HIPAA compliance
206            sqlx::query(
207                r#"
208                CREATE POLICY IF NOT EXISTS hipaa_agent_isolation ON memories
209                FOR ALL
210                USING (agent_id = current_setting('app.current_agent_id')::uuid)
211            "#,
212            )
213            .execute(&self.pool)
214            .await
215            .ok();
216        }
217
218        // Create encrypted column for sensitive data
219        if self.config.encryption_at_rest {
220            // Would use pgcrypto extension
221            sqlx::query("CREATE EXTENSION IF NOT EXISTS pgcrypto")
222                .execute(&self.pool)
223                .await
224                .ok();
225        }
226
227        tracing::info!("HIPAA compliance initialized successfully");
228        Ok(())
229    }
230
231    /// Log an audit entry
232    pub async fn log_audit(&self, entry: AuditLogEntry) -> Result<()> {
233        if !self.config.enabled || !self.config.audit_logging {
234            return Ok(());
235        }
236
237        sqlx::query(
238            "INSERT INTO audit_log (id, timestamp, actor_id, action, resource_type, resource_id, ip_address, result, metadata)
239             VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)"
240        )
241        .bind(entry.id)
242        .bind(entry.timestamp)
243        .bind(entry.actor_id)
244        .bind(&entry.action)
245        .bind(&entry.resource_type)
246        .bind(entry.resource_id)
247        .bind(&entry.ip_address)
248        .bind(&entry.result)
249        .bind(serde_json::to_value(&entry.metadata)?)
250        .execute(&self.pool)
251        .await?;
252
253        Ok(())
254    }
255
256    /// Purge old records per retention policy
257    pub async fn enforce_retention_policy(&self) -> Result<usize> {
258        let cutoff_timestamp =
259            chrono::Utc::now().timestamp() - (self.config.retention_days as i64 * 86400);
260
261        let result = sqlx::query(
262            "DELETE FROM memories WHERE created_at < $1 AND metadata->>'retention' != 'permanent'",
263        )
264        .bind(cutoff_timestamp)
265        .execute(&self.pool)
266        .await?;
267
268        tracing::info!(
269            "Retention policy enforced: {} records purged",
270            result.rows_affected()
271        );
272        Ok(result.rows_affected() as usize)
273    }
274
275    /// Get audit log entries
276    pub async fn get_audit_log(
277        &self,
278        actor_id: Option<uuid::Uuid>,
279        limit: usize,
280    ) -> Result<Vec<AuditLogEntry>> {
281        let query = if let Some(aid) = actor_id {
282            sqlx::query_as::<_, (uuid::Uuid, i64, uuid::Uuid, String, String, uuid::Uuid, Option<String>, String, serde_json::Value)>(
283                "SELECT id, timestamp, actor_id, action, resource_type, resource_id, ip_address, result, metadata 
284                 FROM audit_log WHERE actor_id = $1 ORDER BY timestamp DESC LIMIT $2"
285            )
286            .bind(aid)
287            .bind(limit as i64)
288        } else {
289            sqlx::query_as::<_, (uuid::Uuid, i64, uuid::Uuid, String, String, uuid::Uuid, Option<String>, String, serde_json::Value)>(
290                "SELECT id, timestamp, actor_id, action, resource_type, resource_id, ip_address, result, metadata 
291                 FROM audit_log ORDER BY timestamp DESC LIMIT $1"
292            )
293            .bind(limit as i64)
294        };
295
296        let rows = query.fetch_all(&self.pool).await?;
297
298        let entries = rows
299            .into_iter()
300            .map(
301                |(
302                    id,
303                    timestamp,
304                    actor_id,
305                    action,
306                    resource_type,
307                    resource_id,
308                    ip_address,
309                    result,
310                    metadata,
311                )| {
312                    AuditLogEntry {
313                        id,
314                        timestamp,
315                        actor_id,
316                        action,
317                        resource_type,
318                        resource_id,
319                        ip_address,
320                        result,
321                        metadata: serde_json::from_value(metadata).unwrap_or_default(),
322                    }
323                },
324            )
325            .collect();
326
327        Ok(entries)
328    }
329
330    /// Encrypt Protected Health Information (PHI) using AES-256-GCM
331    ///
332    /// This provides HIPAA-compliant encryption at rest for sensitive medical data.
333    /// The encryption uses:
334    /// - AES-256-GCM for authenticated encryption
335    /// - Argon2 for key derivation
336    /// - Random salt and nonce per encryption
337    ///
338    /// # Errors
339    /// Returns an error if encryption is not properly configured (missing key)
340    pub fn encrypt_phi(&self, data: &str) -> Result<String> {
341        if !self.config.encryption_at_rest {
342            // Encryption disabled - return plaintext (for non-HIPAA deployments)
343            tracing::debug!("PHI encryption disabled, storing plaintext");
344            return Ok(data.to_string());
345        }
346
347        let key = self.config.encryption_key.as_ref().ok_or_else(|| {
348            tracing::error!("CRITICAL: PHI encryption requested but no encryption key configured!");
349            ZoeyError::other(
350                "HIPAA encryption key not configured. Set encryption key via HIPAAConfig::with_encryption_key() \
351                 before encrypting PHI data. This is a HIPAA compliance violation."
352            )
353        })?;
354
355        encrypt_secret(data, key)
356    }
357
358    /// Decrypt Protected Health Information (PHI)
359    ///
360    /// Decrypts data that was encrypted with `encrypt_phi`.
361    ///
362    /// # Errors
363    /// - Returns error if decryption key is not configured
364    /// - Returns error if data was tampered with (authentication failure)
365    /// - Returns error if wrong key is used
366    pub fn decrypt_phi(&self, encrypted: &str) -> Result<String> {
367        if !self.config.encryption_at_rest {
368            // Encryption disabled - data is plaintext
369            return Ok(encrypted.to_string());
370        }
371
372        // Handle legacy "ENCRYPTED:" prefix format (migration path)
373        if encrypted.starts_with("ENCRYPTED:") {
374            tracing::warn!(
375                "Found legacy placeholder encryption format - this data was NOT actually encrypted! \
376                 Re-encrypt this data immediately for HIPAA compliance."
377            );
378            // Return the "encrypted" data for migration purposes, but log a warning
379            return Ok(encrypted
380                .strip_prefix("ENCRYPTED:")
381                .unwrap_or(encrypted)
382                .to_string());
383        }
384
385        let key = self.config.encryption_key.as_ref().ok_or_else(|| {
386            ZoeyError::other("HIPAA encryption key not configured for decryption")
387        })?;
388
389        decrypt_secret(encrypted, key)
390    }
391
392    /// Get the configuration
393    pub fn config(&self) -> &HIPAAConfig {
394        &self.config
395    }
396}
397
398// ============================================================================
399// PII Protection Hooks
400// ============================================================================
401
402/// Access level for protected data
403#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
404pub enum AccessLevel {
405    /// Full access - can see original PII
406    Full,
407    /// Redacted access - can only see redacted version
408    Redacted,
409    /// Metadata only - can see existence but not content
410    MetadataOnly,
411    /// No access
412    None,
413}
414
415/// Document privilege status for access control
416#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
417pub enum DocumentPrivilege {
418    /// Attorney-client privileged
419    AttorneyClient,
420    /// Work product doctrine
421    WorkProduct,
422    /// Both attorney-client and work product
423    Dual,
424    /// Not privileged
425    None,
426}
427
428impl std::fmt::Display for DocumentPrivilege {
429    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
430        match self {
431            DocumentPrivilege::AttorneyClient => write!(f, "Attorney-Client Privileged"),
432            DocumentPrivilege::WorkProduct => write!(f, "Work Product"),
433            DocumentPrivilege::Dual => write!(f, "Privileged (A/C + WP)"),
434            DocumentPrivilege::None => write!(f, "Not Privileged"),
435        }
436    }
437}
438
439/// PII detection result for storage decisions
440#[derive(Debug, Clone, Serialize, Deserialize)]
441pub struct PIIDetectionResult {
442    /// Whether PII was detected
443    pub contains_pii: bool,
444    /// Number of PII instances found
445    pub pii_count: usize,
446    /// Types of PII found
447    pub pii_types: Vec<String>,
448    /// Whether critical PII (SSN, credit card, etc.) was found
449    pub has_critical: bool,
450    /// Suggested access level
451    pub suggested_access: AccessLevel,
452    /// Detected document privilege (for legal documents)
453    pub privilege_status: DocumentPrivilege,
454}
455
456impl Default for PIIDetectionResult {
457    fn default() -> Self {
458        Self {
459            contains_pii: false,
460            pii_count: 0,
461            pii_types: Vec::new(),
462            has_critical: false,
463            suggested_access: AccessLevel::Full,
464            privilege_status: DocumentPrivilege::None,
465        }
466    }
467}
468
469/// Protected data record for dual storage
470#[derive(Debug, Clone, Serialize, Deserialize)]
471pub struct ProtectedDataRecord {
472    /// Unique identifier
473    pub id: uuid::Uuid,
474    /// Original data (encrypted)
475    pub encrypted_original: String,
476    /// Redacted version (safe to display)
477    pub redacted_version: String,
478    /// PII detection metadata
479    pub pii_metadata: PIIDetectionResult,
480    /// Entity ID (owner)
481    pub entity_id: uuid::Uuid,
482    /// Agent ID
483    pub agent_id: uuid::Uuid,
484    /// Room/conversation ID
485    pub room_id: uuid::Uuid,
486    /// Creation timestamp
487    pub created_at: i64,
488    /// Last accessed timestamp
489    pub last_accessed: Option<i64>,
490    /// Access count
491    pub access_count: usize,
492    /// Document type (for legal classification)
493    pub document_type: Option<String>,
494    /// Case reference (for legal documents)
495    pub case_reference: Option<String>,
496}
497
498/// PII Storage Hook for pre-storage processing
499pub struct PIIStorageHook {
500    config: PIIStorageConfig,
501}
502
503/// Configuration for PII storage hooks
504#[derive(Debug, Clone)]
505pub struct PIIStorageConfig {
506    /// Block storage if critical PII is detected (fail-safe mode)
507    pub block_on_critical: bool,
508    /// Always store both encrypted and redacted versions
509    pub dual_storage: bool,
510    /// Automatically detect attorney-client privilege
511    pub detect_privilege: bool,
512    /// Log all PII detections for audit
513    pub audit_detections: bool,
514    /// Redaction placeholder format
515    pub redaction_format: String,
516}
517
518impl Default for PIIStorageConfig {
519    fn default() -> Self {
520        Self {
521            block_on_critical: false, // Don't block by default
522            dual_storage: true,
523            detect_privilege: true,
524            audit_detections: true,
525            redaction_format: "[REDACTED-{}]".to_string(),
526        }
527    }
528}
529
530impl PIIStorageConfig {
531    /// Create config for law office use
532    pub fn for_legal_office() -> Self {
533        Self {
534            block_on_critical: false, // Law offices need to store client data
535            dual_storage: true,
536            detect_privilege: true, // Important for legal
537            audit_detections: true,
538            redaction_format: "[REDACTED-{}]".to_string(),
539        }
540    }
541
542    /// Create config for healthcare/HIPAA use
543    pub fn for_hipaa() -> Self {
544        Self {
545            block_on_critical: true, // Block unintended PHI storage
546            dual_storage: true,
547            detect_privilege: false,
548            audit_detections: true,
549            redaction_format: "[PHI-REDACTED]".to_string(),
550        }
551    }
552
553    /// Create a strict config that blocks on any PII
554    pub fn strict() -> Self {
555        Self {
556            block_on_critical: true,
557            dual_storage: true,
558            detect_privilege: true,
559            audit_detections: true,
560            redaction_format: "[BLOCKED-PII]".to_string(),
561        }
562    }
563}
564
565impl PIIStorageHook {
566    /// Create a new PII storage hook
567    pub fn new(config: PIIStorageConfig) -> Self {
568        Self { config }
569    }
570
571    /// Create with default configuration
572    pub fn default_config() -> Self {
573        Self::new(PIIStorageConfig::default())
574    }
575
576    /// Create for law office use
577    pub fn for_legal_office() -> Self {
578        Self::new(PIIStorageConfig::for_legal_office())
579    }
580
581    /// Create for HIPAA compliance
582    pub fn for_hipaa() -> Self {
583        Self::new(PIIStorageConfig::for_hipaa())
584    }
585
586    /// Pre-storage hook: Analyze text before storage
587    ///
588    /// This performs PII detection and returns metadata about the content.
589    /// The caller can use this to decide how to store the data.
590    pub fn analyze_for_storage(&self, text: &str) -> PIIDetectionResult {
591        let mut result = PIIDetectionResult::default();
592
593        // Basic PII pattern detection (simplified - in production would use HybridPIIDetector)
594        let patterns = [
595            ("SSN", r"\b\d{3}-\d{2}-\d{4}\b"),
596            ("EMAIL", r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,}\b"),
597            ("PHONE", r"\b\d{3}[-.]?\d{3}[-.]?\d{4}\b"),
598            ("CREDITCARD", r"\b\d{4}[\s-]?\d{4}[\s-]?\d{4}[\s-]?\d{4}\b"),
599            ("APIKEY", r"\b(sk-|pk-|api[_-]?key)[A-Za-z0-9_-]+"),
600        ];
601
602        let critical_types = ["SSN", "CREDITCARD", "APIKEY"];
603
604        for (pii_type, pattern) in &patterns {
605            if let Ok(regex) = regex::Regex::new(pattern) {
606                let matches: Vec<_> = regex.find_iter(text).collect();
607                if !matches.is_empty() {
608                    result.contains_pii = true;
609                    result.pii_count += matches.len();
610                    result.pii_types.push(pii_type.to_string());
611
612                    if critical_types.contains(pii_type) {
613                        result.has_critical = true;
614                    }
615                }
616            }
617        }
618
619        // Detect privilege status if enabled
620        if self.config.detect_privilege {
621            result.privilege_status = self.detect_privilege(text);
622        }
623
624        // Determine suggested access level
625        result.suggested_access = if result.has_critical {
626            AccessLevel::Redacted
627        } else if result.contains_pii {
628            AccessLevel::Redacted
629        } else if result.privilege_status != DocumentPrivilege::None {
630            AccessLevel::Redacted // Privileged documents should be protected
631        } else {
632            AccessLevel::Full
633        };
634
635        result
636    }
637
638    /// Detect attorney-client privilege status
639    fn detect_privilege(&self, text: &str) -> DocumentPrivilege {
640        let text_lower = text.to_lowercase();
641
642        let ac_keywords = [
643            "privileged",
644            "attorney-client",
645            "attorney client",
646            "confidential communication",
647            "legal advice",
648            "privileged and confidential",
649        ];
650
651        let wp_keywords = [
652            "work product",
653            "attorney work product",
654            "prepared in anticipation of litigation",
655            "trial preparation",
656            "litigation strategy",
657        ];
658
659        let has_ac = ac_keywords.iter().any(|kw| text_lower.contains(kw));
660        let has_wp = wp_keywords.iter().any(|kw| text_lower.contains(kw));
661
662        match (has_ac, has_wp) {
663            (true, true) => DocumentPrivilege::Dual,
664            (true, false) => DocumentPrivilege::AttorneyClient,
665            (false, true) => DocumentPrivilege::WorkProduct,
666            (false, false) => DocumentPrivilege::None,
667        }
668    }
669
670    /// Redact PII from text using configured format
671    pub fn redact(&self, text: &str) -> String {
672        let mut redacted = text.to_string();
673
674        let patterns = [
675            ("SSN", r"\b\d{3}-\d{2}-\d{4}\b"),
676            ("EMAIL", r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,}\b"),
677            ("PHONE", r"\b\d{3}[-.]?\d{3}[-.]?\d{4}\b"),
678            ("CREDITCARD", r"\b\d{4}[\s-]?\d{4}[\s-]?\d{4}[\s-]?\d{4}\b"),
679            ("APIKEY", r"\b(sk-|pk-|api[_-]?key)[A-Za-z0-9_-]+"),
680        ];
681
682        for (pii_type, pattern) in &patterns {
683            if let Ok(regex) = regex::Regex::new(pattern) {
684                let replacement = self.config.redaction_format.replace("{}", pii_type);
685                redacted = regex.replace_all(&redacted, replacement.as_str()).to_string();
686            }
687        }
688
689        redacted
690    }
691
692    /// Prepare data for dual storage
693    ///
694    /// Returns (encrypted_original, redacted_version)
695    pub fn prepare_for_storage(
696        &self,
697        text: &str,
698        encryption_key: Option<&str>,
699    ) -> Result<(String, String, PIIDetectionResult)> {
700        let detection = self.analyze_for_storage(text);
701
702        // Check if we should block storage
703        if self.config.block_on_critical && detection.has_critical {
704            return Err(ZoeyError::validation(
705                "Storage blocked: Critical PII detected. Please review and redact before storing.",
706            ));
707        }
708
709        // Prepare encrypted original
710        let encrypted = if let Some(key) = encryption_key {
711            encrypt_secret(text, key)?
712        } else {
713            // No encryption key - store plaintext (not recommended for production)
714            text.to_string()
715        };
716
717        // Prepare redacted version
718        let redacted = self.redact(text);
719
720        Ok((encrypted, redacted, detection))
721    }
722
723    /// Check if storage should be allowed based on detection result
724    pub fn should_allow_storage(&self, detection: &PIIDetectionResult) -> bool {
725        if self.config.block_on_critical && detection.has_critical {
726            return false;
727        }
728        true
729    }
730
731    /// Get the access level for a user based on their role
732    pub fn get_access_level(&self, user_role: &str, detection: &PIIDetectionResult) -> AccessLevel {
733        match user_role.to_lowercase().as_str() {
734            "admin" | "attorney" | "partner" => {
735                // Full access for attorneys and partners
736                AccessLevel::Full
737            }
738            "paralegal" | "legal_assistant" => {
739                // Paralegals can see redacted for critical PII
740                if detection.has_critical {
741                    AccessLevel::Redacted
742                } else {
743                    AccessLevel::Full
744                }
745            }
746            "billing" | "accounting" => {
747                // Billing can see metadata and non-critical PII
748                if detection.has_critical {
749                    AccessLevel::MetadataOnly
750                } else {
751                    AccessLevel::Redacted
752                }
753            }
754            "staff" | "receptionist" => {
755                // General staff gets redacted view
756                AccessLevel::Redacted
757            }
758            "external" | "client" => {
759                // External users get metadata only for PII documents
760                if detection.contains_pii {
761                    AccessLevel::MetadataOnly
762                } else {
763                    AccessLevel::Full
764                }
765            }
766            _ => AccessLevel::Redacted, // Default to redacted
767        }
768    }
769}
770
771impl Default for PIIStorageHook {
772    fn default() -> Self {
773        Self::default_config()
774    }
775}
776
777/// Extension trait for HIPAACompliance to add PII storage hooks
778impl HIPAACompliance {
779    /// Create a protected data record with PII detection
780    pub fn create_protected_record(
781        &self,
782        text: &str,
783        entity_id: uuid::Uuid,
784        agent_id: uuid::Uuid,
785        room_id: uuid::Uuid,
786        storage_hook: &PIIStorageHook,
787    ) -> Result<ProtectedDataRecord> {
788        let (encrypted, redacted, detection) = storage_hook.prepare_for_storage(
789            text,
790            self.config.encryption_key.as_deref(),
791        )?;
792
793        Ok(ProtectedDataRecord {
794            id: uuid::Uuid::new_v4(),
795            encrypted_original: encrypted,
796            redacted_version: redacted,
797            pii_metadata: detection,
798            entity_id,
799            agent_id,
800            room_id,
801            created_at: chrono::Utc::now().timestamp(),
802            last_accessed: None,
803            access_count: 0,
804            document_type: None,
805            case_reference: None,
806        })
807    }
808
809    /// Initialize protected data storage table
810    pub async fn initialize_protected_storage(&self) -> Result<()> {
811        if !self.config.enabled {
812            return Ok(());
813        }
814
815        sqlx::query(
816            r#"
817            CREATE TABLE IF NOT EXISTS protected_data (
818                id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
819                encrypted_original TEXT NOT NULL,
820                redacted_version TEXT NOT NULL,
821                pii_metadata JSONB DEFAULT '{}',
822                entity_id UUID NOT NULL,
823                agent_id UUID NOT NULL,
824                room_id UUID NOT NULL,
825                created_at BIGINT NOT NULL,
826                last_accessed BIGINT,
827                access_count INTEGER DEFAULT 0,
828                document_type TEXT,
829                case_reference TEXT,
830                privilege_status TEXT DEFAULT 'none'
831            )
832            "#,
833        )
834        .execute(&self.pool)
835        .await?;
836
837        // Create indexes
838        sqlx::query("CREATE INDEX IF NOT EXISTS idx_protected_data_entity ON protected_data(entity_id)")
839            .execute(&self.pool)
840            .await?;
841        sqlx::query("CREATE INDEX IF NOT EXISTS idx_protected_data_agent ON protected_data(agent_id)")
842            .execute(&self.pool)
843            .await?;
844        sqlx::query("CREATE INDEX IF NOT EXISTS idx_protected_data_case ON protected_data(case_reference)")
845            .execute(&self.pool)
846            .await?;
847
848        // Enable RLS for protected data
849        if self.config.access_control {
850            sqlx::query("ALTER TABLE protected_data ENABLE ROW LEVEL SECURITY")
851                .execute(&self.pool)
852                .await
853                .ok();
854        }
855
856        tracing::info!("Protected data storage initialized");
857        Ok(())
858    }
859
860    /// Store a protected data record
861    pub async fn store_protected_data(&self, record: &ProtectedDataRecord) -> Result<uuid::Uuid> {
862        let pii_metadata_json = serde_json::to_value(&record.pii_metadata)?;
863        let privilege_str = format!("{:?}", record.pii_metadata.privilege_status);
864
865        sqlx::query(
866            r#"
867            INSERT INTO protected_data 
868            (id, encrypted_original, redacted_version, pii_metadata, entity_id, agent_id, room_id, 
869             created_at, last_accessed, access_count, document_type, case_reference, privilege_status)
870            VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13)
871            "#,
872        )
873        .bind(record.id)
874        .bind(&record.encrypted_original)
875        .bind(&record.redacted_version)
876        .bind(pii_metadata_json)
877        .bind(record.entity_id)
878        .bind(record.agent_id)
879        .bind(record.room_id)
880        .bind(record.created_at)
881        .bind(record.last_accessed)
882        .bind(record.access_count as i32)
883        .bind(&record.document_type)
884        .bind(&record.case_reference)
885        .bind(privilege_str)
886        .execute(&self.pool)
887        .await?;
888
889        // Log audit entry
890        if self.config.audit_logging {
891            let mut metadata = HashMap::new();
892            metadata.insert("contains_pii".to_string(), record.pii_metadata.contains_pii.to_string());
893            metadata.insert("pii_count".to_string(), record.pii_metadata.pii_count.to_string());
894            if !record.pii_metadata.pii_types.is_empty() {
895                metadata.insert("pii_types".to_string(), record.pii_metadata.pii_types.join(","));
896            }
897
898            self.log_audit(AuditLogEntry {
899                id: uuid::Uuid::new_v4(),
900                timestamp: chrono::Utc::now().timestamp(),
901                actor_id: record.agent_id,
902                action: "STORE_PROTECTED_DATA".to_string(),
903                resource_type: "ProtectedData".to_string(),
904                resource_id: record.id,
905                ip_address: None,
906                result: "SUCCESS".to_string(),
907                metadata,
908            })
909            .await?;
910        }
911
912        Ok(record.id)
913    }
914
915    /// Retrieve protected data with access control
916    pub async fn get_protected_data(
917        &self,
918        id: uuid::Uuid,
919        accessor_id: uuid::Uuid,
920        access_level: AccessLevel,
921    ) -> Result<Option<(String, PIIDetectionResult)>> {
922        let row = sqlx::query_as::<_, (String, String, serde_json::Value)>(
923            "SELECT encrypted_original, redacted_version, pii_metadata FROM protected_data WHERE id = $1",
924        )
925        .bind(id)
926        .fetch_optional(&self.pool)
927        .await?;
928
929        if let Some((encrypted, redacted, pii_json)) = row {
930            // Update access tracking
931            sqlx::query(
932                "UPDATE protected_data SET last_accessed = $1, access_count = access_count + 1 WHERE id = $2",
933            )
934            .bind(chrono::Utc::now().timestamp())
935            .bind(id)
936            .execute(&self.pool)
937            .await?;
938
939            // Log audit
940            if self.config.audit_logging {
941                let mut metadata = HashMap::new();
942                metadata.insert("access_level".to_string(), format!("{:?}", access_level));
943
944                self.log_audit(AuditLogEntry {
945                    id: uuid::Uuid::new_v4(),
946                    timestamp: chrono::Utc::now().timestamp(),
947                    actor_id: accessor_id,
948                    action: "ACCESS_PROTECTED_DATA".to_string(),
949                    resource_type: "ProtectedData".to_string(),
950                    resource_id: id,
951                    ip_address: None,
952                    result: "SUCCESS".to_string(),
953                    metadata,
954                })
955                .await?;
956            }
957
958            let pii_metadata: PIIDetectionResult = serde_json::from_value(pii_json)?;
959
960            // Return data based on access level
961            let content = match access_level {
962                AccessLevel::Full => {
963                    // Decrypt and return original
964                    self.decrypt_phi(&encrypted)?
965                }
966                AccessLevel::Redacted => redacted,
967                AccessLevel::MetadataOnly => {
968                    format!("[Content protected - {} PII instances detected]", pii_metadata.pii_count)
969                }
970                AccessLevel::None => {
971                    return Err(ZoeyError::auth("Access denied to protected data"));
972                }
973            };
974
975            Ok(Some((content, pii_metadata)))
976        } else {
977            Ok(None)
978        }
979    }
980}
981
982#[cfg(test)]
983mod tests {
984    use super::*;
985
986    #[test]
987    fn test_hipaa_config() {
988        let config = HIPAAConfig::default();
989        assert!(config.audit_logging);
990        assert!(config.encryption_at_rest);
991        assert_eq!(config.retention_days, 2555); // 7 years
992    }
993
994    #[test]
995    fn test_audit_log_entry() {
996        let entry = AuditLogEntry {
997            id: uuid::Uuid::new_v4(),
998            timestamp: chrono::Utc::now().timestamp(),
999            actor_id: uuid::Uuid::new_v4(),
1000            action: "READ".to_string(),
1001            resource_type: "Memory".to_string(),
1002            resource_id: uuid::Uuid::new_v4(),
1003            ip_address: Some("127.0.0.1".to_string()),
1004            result: "SUCCESS".to_string(),
1005            metadata: HashMap::new(),
1006        };
1007
1008        assert_eq!(entry.action, "READ");
1009    }
1010}