lean_ctx/core/

audit_trail.rs

use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use std::fs::{self, OpenOptions};
use std::io::{BufRead, Write};
use std::path::PathBuf;
use std::sync::Mutex;

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AuditEntry {
    pub timestamp: String,
    pub agent_id: String,
    pub tool: String,
    pub action: Option<String>,
    pub input_hash: String,
    pub output_tokens: u32,
    pub role: String,
    pub event_type: AuditEventType,
    pub prev_hash: String,
    pub entry_hash: String,
    /// Ed25519 signature over `entry_hash`, proving provenance from the local
    /// lean-ctx identity. `None` only when the keypair is unavailable (early boot).
    #[serde(skip_serializing_if = "Option::is_none", default)]
    pub signature: Option<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum AuditEventType {
    ToolCall,
    ToolDenied,
    PathJailViolation,
    BudgetExceeded,
    CrossProjectAccess,
    RateLimited,
    SecurityViolation,
    RoleChanged,
    SecretDetected,
}

pub struct AuditEntryData {
    pub agent_id: String,
    pub tool: String,
    pub action: Option<String>,
    pub input_hash: String,
    pub output_tokens: u32,
    pub role: String,
    pub event_type: AuditEventType,
}

pub struct ChainVerifyResult {
    pub total_entries: usize,
    pub valid: bool,
    pub first_invalid_at: Option<usize>,
}

static LAST_HASH: Mutex<Option<String>> = Mutex::new(None);

fn trail_path() -> Option<PathBuf> {
    let dir = crate::core::data_dir::lean_ctx_data_dir().ok()?;
    let audit_dir = dir.join("audit");
    fs::create_dir_all(&audit_dir).ok()?;
    Some(audit_dir.join("trail.jsonl"))
}

fn init_last_hash(path: &PathBuf) -> String {
    let mut guard = LAST_HASH
        .lock()
        .unwrap_or_else(std::sync::PoisonError::into_inner);
    if let Some(ref h) = *guard {
        return h.clone();
    }
    let hash = read_last_hash_from_file(path);
    *guard = Some(hash.clone());
    hash
}

fn read_last_hash_from_file(path: &PathBuf) -> String {
    let Ok(file) = fs::File::open(path) else {
        return "genesis".to_string();
    };
    let reader = std::io::BufReader::new(file);
    let mut last_hash = "genesis".to_string();
    for line in reader.lines().map_while(Result::ok) {
        if let Ok(entry) = serde_json::from_str::<AuditEntry>(&line) {
            last_hash = entry.entry_hash;
        }
    }
    last_hash
}

fn compute_entry_hash(prev_hash: &str, data_json: &str) -> String {
    let mut hasher = Sha256::new();
    hasher.update(prev_hash.as_bytes());
    hasher.update(data_json.as_bytes());
    format!("{:x}", hasher.finalize())
}

pub fn record(data: AuditEntryData) {
    let Some(path) = trail_path() else { return };
    let prev_hash = init_last_hash(&path);

    let partial = serde_json::json!({
        "agent_id": data.agent_id,
        "tool": data.tool,
        "action": data.action,
        "input_hash": data.input_hash,
        "output_tokens": data.output_tokens,
        "role": data.role,
        "event_type": data.event_type,
    });
    let data_json = serde_json::to_string(&partial).unwrap_or_default();
    let entry_hash = compute_entry_hash(&prev_hash, &data_json);

    let signature = crate::core::agent_identity::sign_bytes("lean-ctx", entry_hash.as_bytes())
        .map(|sig| crate::core::agent_identity::hex_encode(&sig))
        .ok();

    let entry = AuditEntry {
        timestamp: chrono::Utc::now().to_rfc3339(),
        agent_id: data.agent_id,
        tool: data.tool,
        action: data.action,
        input_hash: data.input_hash,
        output_tokens: data.output_tokens,
        role: data.role,
        event_type: data.event_type,
        prev_hash,
        entry_hash: entry_hash.clone(),
        signature,
    };

    if let Ok(line) = serde_json::to_string(&entry) {
        if let Ok(mut file) = OpenOptions::new().create(true).append(true).open(&path) {
            let _ = writeln!(file, "{line}");
        }
    }

    if let Ok(mut guard) = LAST_HASH.lock() {
        *guard = Some(entry_hash);
    }
}

pub fn load_recent(limit: usize) -> Vec<AuditEntry> {
    let Some(path) = trail_path() else {
        return Vec::new();
    };
    let Ok(file) = fs::File::open(&path) else {
        return Vec::new();
    };
    let reader = std::io::BufReader::new(file);
    let entries: Vec<AuditEntry> = reader
        .lines()
        .map_while(Result::ok)
        .filter_map(|line| serde_json::from_str(&line).ok())
        .collect();
    let skip = entries.len().saturating_sub(limit);
    entries.into_iter().skip(skip).collect()
}

pub fn verify_chain() -> ChainVerifyResult {
    let Some(path) = trail_path() else {
        return ChainVerifyResult {
            total_entries: 0,
            valid: true,
            first_invalid_at: None,
        };
    };
    let Ok(file) = fs::File::open(&path) else {
        return ChainVerifyResult {
            total_entries: 0,
            valid: true,
            first_invalid_at: None,
        };
    };
    let reader = std::io::BufReader::new(file);
    let mut prev_hash = "genesis".to_string();
    let mut total = 0usize;

    for line in reader.lines().map_while(Result::ok) {
        let entry: AuditEntry = match serde_json::from_str(&line) {
            Ok(e) => e,
            Err(_) => {
                return ChainVerifyResult {
                    total_entries: total,
                    valid: false,
                    first_invalid_at: Some(total),
                }
            }
        };

        if entry.prev_hash != prev_hash {
            return ChainVerifyResult {
                total_entries: total,
                valid: false,
                first_invalid_at: Some(total),
            };
        }

        let partial = serde_json::json!({
            "agent_id": entry.agent_id,
            "tool": entry.tool,
            "action": entry.action,
            "input_hash": entry.input_hash,
            "output_tokens": entry.output_tokens,
            "role": entry.role,
            "event_type": entry.event_type,
        });
        let data_json = serde_json::to_string(&partial).unwrap_or_default();
        let expected = compute_entry_hash(&prev_hash, &data_json);

        if entry.entry_hash != expected {
            return ChainVerifyResult {
                total_entries: total,
                valid: false,
                first_invalid_at: Some(total),
            };
        }

        prev_hash = entry.entry_hash;
        total += 1;
    }

    ChainVerifyResult {
        total_entries: total,
        valid: true,
        first_invalid_at: None,
    }
}

pub fn hash_input(args: &serde_json::Map<String, serde_json::Value>) -> String {
    let serialized = serde_json::to_string(args).unwrap_or_default();
    let mut hasher = Sha256::new();
    hasher.update(serialized.as_bytes());
    format!("{:x}", hasher.finalize())
}