use hmac::{Hmac, Mac};
use serde::{Deserialize, Serialize};
use sha2::Sha256;
use std::path::PathBuf;
use std::sync::Mutex;
use thiserror::Error;
use tracing::info;
use zeroize::Zeroize;
#[allow(dead_code)]
#[derive(Error, Debug)]
pub enum AuditError {
#[error("IO error: {0}")]
Io(#[from] std::io::Error),
#[error("Serialization error: {0}")]
Serialization(#[from] serde_json::Error),
#[error("HMAC verification failed at entry {0}")]
TamperDetected(u64),
#[error("Audit log is empty")]
EmptyLog,
#[error("Audit log not initialized")]
NotInitialized,
}
#[allow(dead_code)]
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[serde(rename_all = "snake_case")]
pub enum AuditEventType {
ToolCall,
ToolResult,
PolicyDecision,
ApprovalRequested,
ApprovalGranted,
ApprovalDenied,
SandboxViolation,
SecurityViolation,
AgentStart,
AgentFinish,
Error,
ConfigChange,
Custom(String),
}
#[allow(dead_code)]
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AuditEntry {
pub index: u64,
pub timestamp: String,
pub event_type: AuditEventType,
pub component: String,
pub description: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub metadata: Option<serde_json::Value>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub session_id: Option<String>,
pub hmac: String,
}
#[allow(dead_code)]
pub struct AuditLog {
key: Vec<u8>,
entries: Mutex<Vec<AuditEntry>>,
session_id: String,
trace_logging: bool,
file_path: Option<PathBuf>,
}
impl Drop for AuditLog {
fn drop(&mut self) {
self.key.zeroize();
}
}
#[allow(dead_code)]
impl AuditLog {
pub fn new(session_id: String) -> Self {
use rand::RngCore;
let mut key = vec![0u8; 32];
rand::rngs::OsRng.fill_bytes(&mut key);
Self {
key,
entries: Mutex::new(Vec::new()),
session_id,
trace_logging: true,
file_path: None,
}
}
pub fn with_key(session_id: String, key: Vec<u8>) -> Self {
Self {
key,
entries: Mutex::new(Vec::new()),
session_id,
trace_logging: true,
file_path: None,
}
}
pub fn set_trace_logging(&mut self, enabled: bool) {
self.trace_logging = enabled;
}
pub fn set_file_path(&mut self, path: PathBuf) {
self.file_path = Some(path);
}
pub fn append(
&self,
event_type: AuditEventType,
component: &str,
description: &str,
metadata: Option<serde_json::Value>,
) -> Result<AuditEntry, AuditError> {
let mut entries = self.lock_entries()?;
let index = entries.len() as u64;
let timestamp = chrono::Utc::now().to_rfc3339();
let hmac = self.compute_hmac(
index,
×tamp,
&event_type,
component,
description,
&metadata,
entries.last(),
)?;
let entry = AuditEntry {
index,
timestamp,
event_type: event_type.clone(),
component: component.to_string(),
description: description.to_string(),
metadata,
session_id: Some(self.session_id.clone()),
hmac,
};
if self.trace_logging {
info!(
audit.index = entry.index,
audit.event = ?entry.event_type,
audit.component = %entry.component,
"Audit: {}",
entry.description
);
}
entries.push(entry.clone());
if let Some(path) = &self.file_path {
self.persist_to_file(path, &entries)?;
}
Ok(entry)
}
pub fn tool_call(
&self,
tool_name: &str,
args: &serde_json::Value,
) -> Result<AuditEntry, AuditError> {
self.append(
AuditEventType::ToolCall,
tool_name,
&format!("Tool call: {}", tool_name),
Some(serde_json::json!({"arguments": args})),
)
}
pub fn tool_result(
&self,
tool_name: &str,
result: &crate::tools::ToolResult,
) -> Result<AuditEntry, AuditError> {
self.append(
AuditEventType::ToolResult,
tool_name,
&format!("Tool result: {} (success: {})", tool_name, result.success),
Some(serde_json::json!({
"success": result.success,
"exit_code": result.exit_code,
"duration_ms": result.duration_ms,
"output_length": result.output.len(),
})),
)
}
pub fn policy_decision(
&self,
tool_name: &str,
allowed: bool,
reason: Option<&str>,
) -> Result<AuditEntry, AuditError> {
self.append(
AuditEventType::PolicyDecision,
"policy",
&format!(
"Policy decision for '{}': {}",
tool_name,
if allowed { "allowed" } else { "denied" }
),
Some(serde_json::json!({
"tool": tool_name,
"allowed": allowed,
"reason": reason,
})),
)
}
pub fn approval(
&self,
tool_name: &str,
granted: bool,
reason: Option<&str>,
) -> Result<AuditEntry, AuditError> {
let event_type = if granted {
AuditEventType::ApprovalGranted
} else {
AuditEventType::ApprovalDenied
};
self.append(
event_type,
"approval",
&format!(
"Approval for '{}': {}",
tool_name,
if granted { "granted" } else { "denied" }
),
Some(serde_json::json!({
"tool": tool_name,
"granted": granted,
"reason": reason,
})),
)
}
pub fn entries(&self) -> Result<Vec<AuditEntry>, AuditError> {
Ok(self.lock_entries()?.clone())
}
pub fn len(&self) -> Result<usize, AuditError> {
Ok(self.lock_entries()?.len())
}
pub fn is_empty(&self) -> Result<bool, AuditError> {
Ok(self.lock_entries()?.is_empty())
}
pub fn verify(&self) -> Result<(), AuditError> {
let entries = self.lock_entries()?;
if entries.is_empty() {
return Err(AuditError::EmptyLog);
}
let mut prev_hmac = String::new();
for entry in entries.iter() {
let expected_hmac = self.compute_hmac_for_verification(
entry.index,
&entry.timestamp,
&entry.event_type,
&entry.component,
&entry.description,
&entry.metadata,
&prev_hmac,
)?;
if entry.hmac != expected_hmac {
return Err(AuditError::TamperDetected(entry.index));
}
prev_hmac = entry.hmac.clone();
}
Ok(())
}
pub fn to_json(&self) -> Result<String, AuditError> {
let entries = self.lock_entries()?;
Ok(serde_json::to_string_pretty(&*entries)?)
}
pub fn to_json_lines(&self) -> Result<String, AuditError> {
let entries = self.lock_entries()?;
let mut lines = String::new();
for entry in entries.iter() {
lines.push_str(&serde_json::to_string(entry)?);
lines.push('\n');
}
Ok(lines)
}
fn lock_entries(&self) -> Result<std::sync::MutexGuard<'_, Vec<AuditEntry>>, AuditError> {
self.entries.lock().map_err(|_| AuditError::NotInitialized)
}
#[allow(clippy::too_many_arguments)]
fn compute_hmac(
&self,
index: u64,
timestamp: &str,
event_type: &AuditEventType,
component: &str,
description: &str,
metadata: &Option<serde_json::Value>,
prev_entry: Option<&AuditEntry>,
) -> Result<String, AuditError> {
let prev_hmac = prev_entry.map(|e| e.hmac.as_str()).unwrap_or("");
self.compute_hmac_for_verification(
index,
timestamp,
event_type,
component,
description,
metadata,
prev_hmac,
)
}
#[allow(clippy::too_many_arguments)]
fn compute_hmac_for_verification(
&self,
index: u64,
timestamp: &str,
event_type: &AuditEventType,
component: &str,
description: &str,
metadata: &Option<serde_json::Value>,
prev_hmac: &str,
) -> Result<String, AuditError> {
let mut mac = Hmac::<Sha256>::new_from_slice(&self.key).map_err(|e| {
AuditError::Serialization(serde_json::Error::io(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
e.to_string(),
)))
})?;
mac.update(prev_hmac.as_bytes());
mac.update(&index.to_le_bytes());
mac.update(timestamp.as_bytes());
mac.update(serde_json::to_string(event_type)?.as_bytes());
mac.update(component.as_bytes());
mac.update(description.as_bytes());
if let Some(m) = metadata {
mac.update(serde_json::to_string(m)?.as_bytes());
}
Ok(hex::encode(mac.finalize().into_bytes()))
}
fn persist_to_file(&self, path: &PathBuf, entries: &[AuditEntry]) -> Result<(), AuditError> {
let json = serde_json::to_string_pretty(entries)?;
std::fs::write(path, json)?;
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
fn create_test_log() -> AuditLog {
AuditLog::with_key("test-session".to_string(), vec![0u8; 32])
}
#[test]
fn test_audit_log_empty() {
let log = create_test_log();
assert!(log.is_empty().unwrap());
assert_eq!(log.len().unwrap(), 0);
}
#[test]
fn test_audit_log_append() {
let log = create_test_log();
let entry = log
.append(AuditEventType::AgentStart, "agent", "Agent started", None)
.unwrap();
assert_eq!(entry.index, 0);
assert_eq!(entry.component, "agent");
assert_eq!(entry.event_type, AuditEventType::AgentStart);
assert!(!entry.hmac.is_empty());
assert_eq!(log.len().unwrap(), 1);
}
#[test]
fn test_audit_log_multiple_entries() {
let log = create_test_log();
log.append(AuditEventType::AgentStart, "agent", "Agent started", None)
.unwrap();
log.append(
AuditEventType::ToolCall,
"shell_exec",
"Executing command",
Some(serde_json::json!({"command": "echo hello"})),
)
.unwrap();
assert_eq!(log.len().unwrap(), 2);
}
#[test]
fn test_audit_log_verify_valid() {
let log = create_test_log();
log.append(AuditEventType::AgentStart, "agent", "Agent started", None)
.unwrap();
log.append(
AuditEventType::ToolCall,
"shell_exec",
"Executing command",
None,
)
.unwrap();
log.append(
AuditEventType::ToolResult,
"shell_exec",
"Command completed",
None,
)
.unwrap();
assert!(log.verify().is_ok());
}
#[test]
fn test_audit_log_verify_tampered() {
let log = create_test_log();
log.append(AuditEventType::AgentStart, "agent", "Agent started", None)
.unwrap();
log.append(
AuditEventType::ToolCall,
"shell_exec",
"Executing command",
None,
)
.unwrap();
{
let mut entries = log.entries.lock().unwrap();
entries[1].description = "Tampered!".to_string();
}
let result = log.verify();
assert!(result.is_err());
assert!(matches!(result.unwrap_err(), AuditError::TamperDetected(1)));
}
#[test]
fn test_audit_log_verify_empty() {
let log = create_test_log();
let result = log.verify();
assert!(matches!(result.unwrap_err(), AuditError::EmptyLog));
}
#[test]
fn test_audit_log_tool_call() {
let log = create_test_log();
let args = serde_json::json!({"command": "echo hello"});
let entry = log.tool_call("shell_exec", &args).unwrap();
assert_eq!(entry.event_type, AuditEventType::ToolCall);
assert_eq!(entry.component, "shell_exec");
}
#[test]
fn test_audit_log_policy_decision() {
let log = create_test_log();
let entry = log.policy_decision("shell_exec", true, None).unwrap();
assert_eq!(entry.event_type, AuditEventType::PolicyDecision);
assert_eq!(entry.component, "policy");
}
#[test]
fn test_audit_log_approval_granted() {
let log = create_test_log();
let entry = log
.approval("shell_exec", true, Some("User approved"))
.unwrap();
assert_eq!(entry.event_type, AuditEventType::ApprovalGranted);
}
#[test]
fn test_audit_log_approval_denied() {
let log = create_test_log();
let entry = log
.approval("shell_exec", false, Some("User denied"))
.unwrap();
assert_eq!(entry.event_type, AuditEventType::ApprovalDenied);
}
#[test]
fn test_audit_log_to_json() {
let log = create_test_log();
log.append(AuditEventType::AgentStart, "agent", "Agent started", None)
.unwrap();
let json = log.to_json().unwrap();
assert!(json.contains("Agent started"));
assert!(json.contains("agent_start"));
}
#[test]
fn test_audit_log_to_json_lines() {
let log = create_test_log();
log.append(AuditEventType::AgentStart, "agent", "Agent started", None)
.unwrap();
log.append(
AuditEventType::ToolCall,
"shell_exec",
"Executing command",
None,
)
.unwrap();
let lines = log.to_json_lines().unwrap();
let line_count = lines.lines().count();
assert_eq!(line_count, 2);
}
#[test]
fn test_audit_log_session_id() {
let log = AuditLog::new("my-session-123".to_string());
let entry = log
.append(AuditEventType::AgentStart, "agent", "Agent started", None)
.unwrap();
assert_eq!(entry.session_id.unwrap(), "my-session-123");
}
#[test]
fn test_audit_log_different_keys_produce_different_hmacs() {
let log1 = AuditLog::with_key("test".to_string(), vec![1u8; 32]);
let log2 = AuditLog::with_key("test".to_string(), vec![2u8; 32]);
let e1 = log1
.append(AuditEventType::AgentStart, "agent", "Agent started", None)
.unwrap();
let e2 = log2
.append(AuditEventType::AgentStart, "agent", "Agent started", None)
.unwrap();
assert_ne!(e1.hmac, e2.hmac);
}
#[test]
fn test_audit_entry_serialization() {
let entry = AuditEntry {
index: 0,
timestamp: "2026-01-01T00:00:00Z".to_string(),
event_type: AuditEventType::ToolCall,
component: "shell_exec".to_string(),
description: "Executed command".to_string(),
metadata: Some(serde_json::json!({"command": "echo hello"})),
session_id: Some("session-1".to_string()),
hmac: "abc123".to_string(),
};
let json = serde_json::to_string(&entry).unwrap();
assert!(json.contains("shell_exec"));
assert!(json.contains("tool_call"));
assert!(json.contains("echo hello"));
}
#[test]
fn test_audit_error_tamper_detected() {
let err = AuditError::TamperDetected(5);
assert_eq!(format!("{}", err), "HMAC verification failed at entry 5");
}
#[test]
fn test_audit_error_empty() {
let err = AuditError::EmptyLog;
assert_eq!(format!("{}", err), "Audit log is empty");
}
#[test]
fn test_audit_event_type_custom() {
let event = AuditEventType::Custom("my_event".to_string());
let json = serde_json::to_string(&event).unwrap();
assert_eq!(json, "{\"custom\":\"my_event\"}");
}
#[test]
fn test_audit_log_with_metadata() {
let log = create_test_log();
let metadata = serde_json::json!({
"key1": "value1",
"key2": 42,
"nested": {"a": 1}
});
let entry = log
.append(
AuditEventType::ConfigChange,
"config",
"Configuration changed",
Some(metadata),
)
.unwrap();
assert!(entry.metadata.is_some());
let meta = entry.metadata.unwrap();
assert_eq!(meta["key1"], "value1");
assert_eq!(meta["key2"], 42);
}
#[test]
fn test_audit_log_trace_logging_disabled() {
let mut log = create_test_log();
log.set_trace_logging(false);
let entry = log
.append(AuditEventType::AgentStart, "agent", "Agent started", None)
.unwrap();
assert_eq!(entry.index, 0);
}
#[test]
fn test_audit_log_chain_integrity() {
let log = create_test_log();
for i in 0..10 {
log.append(
AuditEventType::Custom(format!("event_{}", i)),
"test",
&format!("Entry {}", i),
None,
)
.unwrap();
}
assert!(log.verify().is_ok());
{
let mut entries = log.entries.lock().unwrap();
entries[3].description = "MODIFIED".to_string();
}
assert!(log.verify().is_err());
}
}