tryaudex_core/
dbaudit.rs

1use serde::{Deserialize, Serialize};
2
3use crate::audit::{AuditEntry, AuditEvent};
4use crate::error::{AvError, Result};
5use crate::session::Session;
6
7/// Database audit backend type.
8#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
9#[serde(rename_all = "lowercase")]
10pub enum DbBackend {
11    Sqlite,
12    Postgres,
13}
14
15impl std::fmt::Display for DbBackend {
16    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
17        match self {
18            Self::Sqlite => write!(f, "sqlite"),
19            Self::Postgres => write!(f, "postgres"),
20        }
21    }
22}
23
24/// Configuration for the database audit backend.
25#[derive(Debug, Clone, Serialize, Deserialize)]
26pub struct DbAuditConfig {
27    /// Backend type: "sqlite" or "postgres"
28    pub backend: DbBackend,
29    /// Connection string.
30    /// - SQLite: file path (e.g. "~/.local/share/audex/audit.db") or ":memory:"
31    /// - Postgres: connection URI (e.g. "postgres://user:pass@localhost/audex")
32    pub connection: String,
33    /// Maximum number of connections in the pool (postgres only, default: 5)
34    #[serde(default = "default_max_connections")]
35    pub max_connections: u32,
36    /// Enable WAL mode for SQLite (default: true)
37    #[serde(default = "default_wal_mode")]
38    pub wal_mode: bool,
39    /// Retention period in days (0 = keep forever, default: 0)
40    #[serde(default)]
41    pub retention_days: u32,
42}
43
44fn default_max_connections() -> u32 {
45    5
46}
47
48fn default_wal_mode() -> bool {
49    true
50}
51
52impl Default for DbAuditConfig {
53    fn default() -> Self {
54        let path = dirs::data_local_dir()
55            .unwrap_or_else(|| std::path::PathBuf::from("."))
56            .join("audex")
57            .join("audit.db");
58        Self {
59            backend: DbBackend::Sqlite,
60            connection: path.to_string_lossy().to_string(),
61            max_connections: default_max_connections(),
62            wal_mode: default_wal_mode(),
63            retention_days: 0,
64        }
65    }
66}
67
68/// SQL schema for audit tables.
69pub const SQLITE_SCHEMA: &str = r#"
70CREATE TABLE IF NOT EXISTS audit_entries (
71    id INTEGER PRIMARY KEY AUTOINCREMENT,
72    timestamp TEXT NOT NULL,
73    session_id TEXT NOT NULL,
74    provider TEXT NOT NULL DEFAULT 'aws',
75    event_type TEXT NOT NULL,
76    event_data TEXT NOT NULL,
77    created_at TEXT NOT NULL DEFAULT (datetime('now'))
78);
79
80CREATE INDEX IF NOT EXISTS idx_audit_session_id ON audit_entries(session_id);
81CREATE INDEX IF NOT EXISTS idx_audit_timestamp ON audit_entries(timestamp);
82CREATE INDEX IF NOT EXISTS idx_audit_provider ON audit_entries(provider);
83CREATE INDEX IF NOT EXISTS idx_audit_event_type ON audit_entries(event_type);
84"#;
85
86pub const POSTGRES_SCHEMA: &str = r#"
87CREATE TABLE IF NOT EXISTS audit_entries (
88    id BIGSERIAL PRIMARY KEY,
89    timestamp TIMESTAMPTZ NOT NULL,
90    session_id TEXT NOT NULL,
91    provider TEXT NOT NULL DEFAULT 'aws',
92    event_type TEXT NOT NULL,
93    event_data JSONB NOT NULL,
94    created_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
95);
96
97CREATE INDEX IF NOT EXISTS idx_audit_session_id ON audit_entries(session_id);
98CREATE INDEX IF NOT EXISTS idx_audit_timestamp ON audit_entries(timestamp);
99CREATE INDEX IF NOT EXISTS idx_audit_provider ON audit_entries(provider);
100CREATE INDEX IF NOT EXISTS idx_audit_event_type ON audit_entries(event_type);
101"#;
102
103/// Extract event type string from an AuditEvent.
104pub fn event_type_str(event: &AuditEvent) -> &'static str {
105    match event {
106        AuditEvent::SessionCreated { .. } => "session_created",
107        AuditEvent::CredentialsIssued { .. } => "credentials_issued",
108        AuditEvent::SessionEnded { .. } => "session_ended",
109        AuditEvent::BudgetWarning { .. } => "budget_warning",
110        AuditEvent::BudgetExceeded { .. } => "budget_exceeded",
111    }
112}
113
114/// SQL statements for common operations.
115pub struct SqlStatements;
116
117impl SqlStatements {
118    /// Insert an audit entry.
119    pub fn insert(backend: &DbBackend) -> &'static str {
120        match backend {
121            DbBackend::Sqlite => {
122                "INSERT INTO audit_entries (timestamp, session_id, provider, event_type, event_data) VALUES (?, ?, ?, ?, ?)"
123            }
124            DbBackend::Postgres => {
125                "INSERT INTO audit_entries (timestamp, session_id, provider, event_type, event_data) VALUES ($1, $2, $3, $4, $5::jsonb)"
126            }
127        }
128    }
129
130    /// Query entries by session ID.
131    pub fn by_session(backend: &DbBackend) -> &'static str {
132        match backend {
133            DbBackend::Sqlite => {
134                "SELECT timestamp, session_id, provider, event_type, event_data FROM audit_entries WHERE session_id = ? ORDER BY timestamp ASC"
135            }
136            DbBackend::Postgres => {
137                "SELECT timestamp, session_id, provider, event_type, event_data FROM audit_entries WHERE session_id = $1 ORDER BY timestamp ASC"
138            }
139        }
140    }
141
142    /// Query recent entries with optional limit.
143    pub fn recent(backend: &DbBackend) -> &'static str {
144        match backend {
145            DbBackend::Sqlite => {
146                "SELECT timestamp, session_id, provider, event_type, event_data FROM audit_entries ORDER BY timestamp DESC LIMIT ?"
147            }
148            DbBackend::Postgres => {
149                "SELECT timestamp, session_id, provider, event_type, event_data FROM audit_entries ORDER BY timestamp DESC LIMIT $1"
150            }
151        }
152    }
153
154    /// Query entries by provider.
155    pub fn by_provider(backend: &DbBackend) -> &'static str {
156        match backend {
157            DbBackend::Sqlite => {
158                "SELECT timestamp, session_id, provider, event_type, event_data FROM audit_entries WHERE provider = ? ORDER BY timestamp DESC LIMIT ?"
159            }
160            DbBackend::Postgres => {
161                "SELECT timestamp, session_id, provider, event_type, event_data FROM audit_entries WHERE provider = $1 ORDER BY timestamp DESC LIMIT $2"
162            }
163        }
164    }
165
166    /// Count entries.
167    pub fn count(backend: &DbBackend) -> &'static str {
168        match backend {
169            DbBackend::Sqlite => "SELECT COUNT(*) FROM audit_entries",
170            DbBackend::Postgres => "SELECT COUNT(*) FROM audit_entries",
171        }
172    }
173
174    /// Count entries by event type.
175    pub fn count_by_type(backend: &DbBackend) -> &'static str {
176        match backend {
177            DbBackend::Sqlite => {
178                "SELECT event_type, COUNT(*) as count FROM audit_entries GROUP BY event_type ORDER BY count DESC"
179            }
180            DbBackend::Postgres => {
181                "SELECT event_type, COUNT(*) as count FROM audit_entries GROUP BY event_type ORDER BY count DESC"
182            }
183        }
184    }
185
186    /// Delete entries older than N days.
187    pub fn prune(backend: &DbBackend) -> &'static str {
188        match backend {
189            DbBackend::Sqlite => {
190                "DELETE FROM audit_entries WHERE timestamp < datetime('now', '-' || ? || ' days')"
191            }
192            DbBackend::Postgres => {
193                "DELETE FROM audit_entries WHERE timestamp < NOW() - ($1 || ' days')::INTERVAL"
194            }
195        }
196    }
197
198    /// Search event data with text matching.
199    pub fn search(backend: &DbBackend) -> &'static str {
200        match backend {
201            DbBackend::Sqlite => {
202                "SELECT timestamp, session_id, provider, event_type, event_data FROM audit_entries WHERE event_data LIKE ? ORDER BY timestamp DESC LIMIT ?"
203            }
204            DbBackend::Postgres => {
205                "SELECT timestamp, session_id, provider, event_type, event_data FROM audit_entries WHERE event_data::text ILIKE $1 ORDER BY timestamp DESC LIMIT $2"
206            }
207        }
208    }
209}
210
211/// Convert an AuditEntry to database row values.
212pub fn entry_to_row(entry: &AuditEntry) -> Result<DbRow> {
213    let event_data = serde_json::to_string(&entry.event)
214        .map_err(|e| AvError::Sts(format!("Failed to serialize event: {}", e)))?;
215
216    // Redact secrets before storing
217    let event_data = crate::leakdetect::redact_secrets(&event_data);
218
219    Ok(DbRow {
220        timestamp: entry.timestamp.to_rfc3339(),
221        session_id: entry.session_id.clone(),
222        provider: entry.provider.clone(),
223        event_type: event_type_str(&entry.event).to_string(),
224        event_data,
225    })
226}
227
228/// Convert a Session to an AuditEntry for session_created events.
229pub fn session_created_entry(session: &Session) -> AuditEntry {
230    use crate::session::CloudProvider;
231
232    let allowed_actions: Vec<String> = match session.provider {
233        CloudProvider::Gcp => session
234            .policy
235            .actions
236            .iter()
237            .map(|a| a.to_gcp_permission())
238            .collect(),
239        CloudProvider::Azure => session
240            .policy
241            .actions
242            .iter()
243            .map(|a| a.to_azure_permission())
244            .collect(),
245        CloudProvider::Aws => session
246            .policy
247            .actions
248            .iter()
249            .map(|a| a.to_iam_action())
250            .collect(),
251    };
252
253    AuditEntry {
254        timestamp: chrono::Utc::now(),
255        session_id: session.id.clone(),
256        provider: session.provider.to_string(),
257        event: AuditEvent::SessionCreated {
258            role_arn: session.role_arn.clone(),
259            ttl_seconds: session.ttl_seconds,
260            budget: session.budget,
261            allowed_actions,
262            command: session.command.clone(),
263            agent_id: session.agent_id.clone(),
264        },
265    }
266}
267
268/// Database row representation.
269#[derive(Debug, Clone, Serialize, Deserialize)]
270pub struct DbRow {
271    pub timestamp: String,
272    pub session_id: String,
273    pub provider: String,
274    pub event_type: String,
275    pub event_data: String,
276}
277
278impl DbRow {
279    /// Convert back to an AuditEntry.
280    pub fn to_entry(&self) -> Result<AuditEntry> {
281        let timestamp = chrono::DateTime::parse_from_rfc3339(&self.timestamp)
282            .map(|dt| dt.with_timezone(&chrono::Utc))
283            .map_err(|e| AvError::Sts(format!("Invalid timestamp: {}", e)))?;
284
285        let event: AuditEvent = serde_json::from_str(&self.event_data)
286            .map_err(|e| AvError::Sts(format!("Invalid event data: {}", e)))?;
287
288        Ok(AuditEntry {
289            timestamp,
290            session_id: self.session_id.clone(),
291            provider: self.provider.clone(),
292            event,
293        })
294    }
295}
296
297/// Statistics about the audit database.
298#[derive(Debug, Clone, Serialize, Deserialize)]
299pub struct DbAuditStats {
300    pub total_entries: u64,
301    pub entries_by_type: Vec<(String, u64)>,
302    pub backend: String,
303    pub connection: String,
304}
305
306/// Migration helper — import JSONL audit log into database.
307pub fn parse_jsonl_for_import(content: &str) -> Vec<AuditEntry> {
308    content
309        .lines()
310        .filter(|l| !l.trim().is_empty())
311        .filter_map(|l| {
312            let (json_content, _) = crate::integrity::parse_line(l);
313            serde_json::from_str(json_content).ok()
314        })
315        .collect()
316}
317
318/// Generate a migration script to import existing JSONL data.
319pub fn generate_import_sql(entries: &[AuditEntry], backend: &DbBackend) -> Result<String> {
320    let mut sql = String::new();
321
322    match backend {
323        DbBackend::Sqlite => {
324            sql.push_str("BEGIN TRANSACTION;\n");
325        }
326        DbBackend::Postgres => {
327            sql.push_str("BEGIN;\n");
328        }
329    }
330
331    for entry in entries {
332        let row = entry_to_row(entry)?;
333        let escaped_data = row.event_data.replace('\'', "''");
334        sql.push_str(&format!(
335            "INSERT INTO audit_entries (timestamp, session_id, provider, event_type, event_data) VALUES ('{}', '{}', '{}', '{}', '{}');\n",
336            row.timestamp, row.session_id, row.provider, row.event_type, escaped_data
337        ));
338    }
339
340    sql.push_str("COMMIT;\n");
341    Ok(sql)
342}
343
344#[cfg(test)]
345mod tests {
346    use super::*;
347    use crate::audit::AuditEvent;
348
349    #[test]
350    fn test_config_default() {
351        let config = DbAuditConfig::default();
352        assert_eq!(config.backend, DbBackend::Sqlite);
353        assert!(config.connection.contains("audit.db"));
354        assert_eq!(config.max_connections, 5);
355        assert!(config.wal_mode);
356    }
357
358    #[test]
359    fn test_config_deserialize_sqlite() {
360        let toml_str = r#"
361backend = "sqlite"
362connection = "/data/audex/audit.db"
363wal_mode = true
364retention_days = 90
365"#;
366        let config: DbAuditConfig = toml::from_str(toml_str).unwrap();
367        assert_eq!(config.backend, DbBackend::Sqlite);
368        assert_eq!(config.connection, "/data/audex/audit.db");
369        assert_eq!(config.retention_days, 90);
370    }
371
372    #[test]
373    fn test_config_deserialize_postgres() {
374        let toml_str = r#"
375backend = "postgres"
376connection = "postgres://audex:password@db.internal:5432/audex"
377max_connections = 10
378"#;
379        let config: DbAuditConfig = toml::from_str(toml_str).unwrap();
380        assert_eq!(config.backend, DbBackend::Postgres);
381        assert_eq!(config.max_connections, 10);
382    }
383
384    #[test]
385    fn test_event_type_str() {
386        assert_eq!(
387            event_type_str(&AuditEvent::SessionCreated {
388                role_arn: "arn".into(),
389                ttl_seconds: 900,
390                budget: None,
391                allowed_actions: vec![],
392                command: vec![],
393                agent_id: None,
394            }),
395            "session_created"
396        );
397        assert_eq!(
398            event_type_str(&AuditEvent::SessionEnded {
399                status: "ok".into(),
400                duration_seconds: 60,
401                exit_code: Some(0),
402            }),
403            "session_ended"
404        );
405        assert_eq!(
406            event_type_str(&AuditEvent::BudgetWarning {
407                current_spend: 1.0,
408                limit: 5.0,
409            }),
410            "budget_warning"
411        );
412    }
413
414    #[test]
415    fn test_entry_to_row() {
416        let entry = AuditEntry {
417            timestamp: chrono::Utc::now(),
418            session_id: "test-123".to_string(),
419            provider: "aws".to_string(),
420            event: AuditEvent::SessionCreated {
421                role_arn: "arn:aws:iam::123:role/Test".to_string(),
422                ttl_seconds: 900,
423                budget: Some(5.0),
424                allowed_actions: vec!["s3:GetObject".to_string()],
425                command: vec!["aws".to_string(), "s3".to_string(), "ls".to_string()],
426                agent_id: Some("claude".to_string()),
427            },
428        };
429        let row = entry_to_row(&entry).unwrap();
430        assert_eq!(row.session_id, "test-123");
431        assert_eq!(row.provider, "aws");
432        assert_eq!(row.event_type, "session_created");
433        assert!(row.event_data.contains("s3:GetObject"));
434    }
435
436    #[test]
437    fn test_row_roundtrip() {
438        let entry = AuditEntry {
439            timestamp: chrono::Utc::now(),
440            session_id: "rt-456".to_string(),
441            provider: "gcp".to_string(),
442            event: AuditEvent::SessionEnded {
443                status: "completed".to_string(),
444                duration_seconds: 120,
445                exit_code: Some(0),
446            },
447        };
448        let row = entry_to_row(&entry).unwrap();
449        let restored = row.to_entry().unwrap();
450        assert_eq!(restored.session_id, "rt-456");
451        assert_eq!(restored.provider, "gcp");
452    }
453
454    #[test]
455    fn test_sql_statements_sqlite() {
456        let b = DbBackend::Sqlite;
457        assert!(SqlStatements::insert(&b).contains("?"));
458        assert!(SqlStatements::by_session(&b).contains("session_id = ?"));
459        assert!(SqlStatements::recent(&b).contains("LIMIT ?"));
460        assert!(SqlStatements::prune(&b).contains("datetime"));
461    }
462
463    #[test]
464    fn test_sql_statements_postgres() {
465        let b = DbBackend::Postgres;
466        assert!(SqlStatements::insert(&b).contains("$1"));
467        assert!(SqlStatements::by_session(&b).contains("session_id = $1"));
468        assert!(SqlStatements::recent(&b).contains("LIMIT $1"));
469        assert!(SqlStatements::prune(&b).contains("INTERVAL"));
470    }
471
472    #[test]
473    fn test_parse_jsonl_for_import() {
474        let jsonl = r#"{"timestamp":"2026-01-01T00:00:00Z","session_id":"s1","provider":"aws","event":{"type":"session_ended","status":"ok","duration_seconds":60,"exit_code":0}}"#;
475        let entries = parse_jsonl_for_import(jsonl);
476        assert_eq!(entries.len(), 1);
477        assert_eq!(entries[0].session_id, "s1");
478    }
479
480    #[test]
481    fn test_generate_import_sql_sqlite() {
482        let entries = vec![AuditEntry {
483            timestamp: chrono::DateTime::parse_from_rfc3339("2026-01-01T00:00:00Z")
484                .unwrap()
485                .with_timezone(&chrono::Utc),
486            session_id: "s1".to_string(),
487            provider: "aws".to_string(),
488            event: AuditEvent::SessionEnded {
489                status: "ok".to_string(),
490                duration_seconds: 60,
491                exit_code: Some(0),
492            },
493        }];
494        let sql = generate_import_sql(&entries, &DbBackend::Sqlite).unwrap();
495        assert!(sql.contains("BEGIN TRANSACTION"));
496        assert!(sql.contains("INSERT INTO audit_entries"));
497        assert!(sql.contains("COMMIT"));
498    }
499
500    #[test]
501    fn test_generate_import_sql_postgres() {
502        let entries = vec![AuditEntry {
503            timestamp: chrono::DateTime::parse_from_rfc3339("2026-01-01T00:00:00Z")
504                .unwrap()
505                .with_timezone(&chrono::Utc),
506            session_id: "s2".to_string(),
507            provider: "gcp".to_string(),
508            event: AuditEvent::CredentialsIssued {
509                access_key_id: "AKID".to_string(),
510                expires_at: chrono::DateTime::parse_from_rfc3339("2026-01-01T00:15:00Z")
511                    .unwrap()
512                    .with_timezone(&chrono::Utc),
513            },
514        }];
515        let sql = generate_import_sql(&entries, &DbBackend::Postgres).unwrap();
516        assert!(sql.contains("BEGIN;"));
517        assert!(sql.contains("COMMIT;"));
518    }
519
520    #[test]
521    fn test_db_backend_display() {
522        assert_eq!(DbBackend::Sqlite.to_string(), "sqlite");
523        assert_eq!(DbBackend::Postgres.to_string(), "postgres");
524    }
525
526    #[test]
527    fn test_db_audit_stats_serialization() {
528        let stats = DbAuditStats {
529            total_entries: 1000,
530            entries_by_type: vec![
531                ("session_created".to_string(), 400),
532                ("credentials_issued".to_string(), 400),
533                ("session_ended".to_string(), 200),
534            ],
535            backend: "sqlite".to_string(),
536            connection: "/data/audit.db".to_string(),
537        };
538        let json = serde_json::to_string(&stats).unwrap();
539        assert!(json.contains("1000"));
540        assert!(json.contains("session_created"));
541    }
542}
tryaudex_core/dbaudit.rs

tryaudex_core/
dbaudit.rs
