use crate::core::config::Config;
use crate::core::db::DbManager;
use crate::core::diff_engine::compute_diff;
use crate::core::policy::{PolicyAction, PolicyEngine};
use crate::core::snapshot::{compute_changes, Change, ChangeType, Snapshot};
use crate::models::{Actor, ActorType, Event, Payload};
use anyhow::Result;
use colored::Colorize;
use std::path::Path;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use tokio::sync::Mutex;
pub struct SessionCapture {
session_id: String,
agent_name: String,
db: Arc<Mutex<DbManager>>,
baseline: Arc<Mutex<Option<Snapshot>>>,
running: Arc<AtomicBool>,
config: Config,
policy_engine: Option<PolicyEngine>,
violations: Arc<Mutex<Vec<String>>>,
}
impl SessionCapture {
pub fn new<P: AsRef<Path>>(
session_id: String,
agent_name: String,
db_path: P,
config: Config,
) -> Result<Self> {
let db = DbManager::open(db_path)?;
let policy_engine = if config.policy.enabled {
match PolicyEngine::new(config.policy.clone()) {
Ok(engine) => Some(engine),
Err(e) => {
eprintln!("Warning: Failed to initialize policy engine: {}", e);
None
}
}
} else {
None
};
Ok(Self {
session_id,
agent_name,
db: Arc::new(Mutex::new(db)),
baseline: Arc::new(Mutex::new(None)),
running: Arc::new(AtomicBool::new(true)),
config,
policy_engine,
violations: Arc::new(Mutex::new(Vec::new())),
})
}
pub async fn capture_baseline<P: AsRef<Path>>(&self, root: P) -> Result<()> {
let snapshot = Snapshot::capture(&root, &self.config.settings.ignore_patterns)?;
let mut baseline = self.baseline.lock().await;
*baseline = Some(snapshot);
Ok(())
}
pub async fn detect_and_record_changes<P: AsRef<Path>>(&self, root: P) -> Result<usize> {
let new_snapshot = Snapshot::capture(&root, &self.config.settings.ignore_patterns)?;
let baseline = self.baseline.lock().await;
let changes = if let Some(ref old) = *baseline {
compute_changes(old, &new_snapshot)
} else {
Vec::new()
};
drop(baseline);
let mut allowed_changes = Vec::new();
let mut blocked_count = 0;
for change in changes {
let path_str = change.path.to_string_lossy().to_string();
match self.check_policy(&change.path).await {
PolicyCheckResult::Allow => {
allowed_changes.push(change);
}
PolicyCheckResult::Block => {
blocked_count += 1;
let mut violations = self.violations.lock().await;
violations.push(format!("BLOCKED: {}", path_str));
}
PolicyCheckResult::Warn => {
allowed_changes.push(change);
println!(
"{}",
format!("⚠️ WARNING: Modifying {}", path_str).yellow()
);
}
PolicyCheckResult::RequireConfirmation => {
allowed_changes.push(change);
println!(
"{}",
format!("⚠️ CONFIRMATION REQUIRED: {}", path_str).cyan()
);
}
}
}
if blocked_count > 0 {
println!(
"{}",
format!(
"🚫 Blocked {} file change(s) due to policy violations",
blocked_count
)
.red()
.bold()
);
}
let change_count = allowed_changes.len();
for change in allowed_changes {
self.record_change(change).await?;
}
let mut baseline = self.baseline.lock().await;
*baseline = Some(new_snapshot);
Ok(change_count)
}
async fn check_policy(&self, path: &std::path::Path) -> PolicyCheckResult {
if let Some(ref engine) = self.policy_engine {
if let Some(violation) = engine.check_file(path) {
match violation.action {
PolicyAction::Block => PolicyCheckResult::Block,
PolicyAction::Warn => PolicyCheckResult::Warn,
PolicyAction::RequireConfirmation => PolicyCheckResult::RequireConfirmation,
PolicyAction::Allow => PolicyCheckResult::Allow,
}
} else {
PolicyCheckResult::Allow
}
} else {
PolicyCheckResult::Allow
}
}
async fn record_change(&self, change: Change) -> Result<()> {
let (lines_added, lines_removed, diff) = match change.change_type {
ChangeType::Created => {
let content = change.new_content.as_deref().unwrap_or("");
let lines = content.lines().count();
(lines, 0, format!("+{}", content))
}
ChangeType::Deleted => {
let content = change.old_content.as_deref().unwrap_or("");
let lines = content.lines().count();
(0, lines, format!("-{}", content))
}
ChangeType::Modified => {
let old = change.old_content.as_deref().unwrap_or("");
let new = change.new_content.as_deref().unwrap_or("");
let result = compute_diff(old, new);
(result.lines_added, result.lines_removed, result.diff_text)
}
};
let diff_to_store = if self.config.settings.capture_diffs {
Some(diff)
} else {
None
};
let event = Event::new(
self.session_id.clone(),
crate::models::EventType::FileChange,
Actor {
r#type: ActorType::Ai,
agent_name: Some(self.agent_name.clone()),
},
Payload::FileChange {
file: change.path.to_string_lossy().to_string(),
lines_added,
lines_removed,
diff: diff_to_store,
},
);
let db = self.db.lock().await;
db.insert_event(&event)?;
Ok(())
}
pub fn is_running(&self) -> bool {
self.running.load(Ordering::Relaxed)
}
pub fn stop(&self) {
self.running.store(false, Ordering::Relaxed);
}
pub async fn get_violations(&self) -> Vec<String> {
self.violations.lock().await.clone()
}
}
enum PolicyCheckResult {
Allow,
Block,
Warn,
RequireConfirmation,
}