1use serde::{Deserialize, Serialize};
2use sha2::{Digest, Sha256};
3use std::fs::{self, OpenOptions};
4use std::io::{BufRead, Write};
5use std::path::PathBuf;
6use std::sync::Mutex;
7
8#[derive(Debug, Clone, Serialize, Deserialize)]
9pub struct AuditEntry {
10 pub timestamp: String,
11 pub agent_id: String,
12 pub tool: String,
13 pub action: Option<String>,
14 pub input_hash: String,
15 pub output_tokens: u32,
16 pub role: String,
17 pub event_type: AuditEventType,
18 pub prev_hash: String,
19 pub entry_hash: String,
20}
21
22#[derive(Debug, Clone, Serialize, Deserialize)]
23#[serde(rename_all = "snake_case")]
24pub enum AuditEventType {
25 ToolCall,
26 ToolDenied,
27 PathJailViolation,
28 BudgetExceeded,
29 CrossProjectAccess,
30 RateLimited,
31 SecurityViolation,
32 RoleChanged,
33 SecretDetected,
34}
35
36pub struct AuditEntryData {
37 pub agent_id: String,
38 pub tool: String,
39 pub action: Option<String>,
40 pub input_hash: String,
41 pub output_tokens: u32,
42 pub role: String,
43 pub event_type: AuditEventType,
44}
45
46pub struct ChainVerifyResult {
47 pub total_entries: usize,
48 pub valid: bool,
49 pub first_invalid_at: Option<usize>,
50}
51
52static LAST_HASH: Mutex<Option<String>> = Mutex::new(None);
53
54fn trail_path() -> Option<PathBuf> {
55 let dir = crate::core::data_dir::lean_ctx_data_dir().ok()?;
56 let audit_dir = dir.join("audit");
57 fs::create_dir_all(&audit_dir).ok()?;
58 Some(audit_dir.join("trail.jsonl"))
59}
60
61fn init_last_hash(path: &PathBuf) -> String {
62 let mut guard = LAST_HASH
63 .lock()
64 .unwrap_or_else(std::sync::PoisonError::into_inner);
65 if let Some(ref h) = *guard {
66 return h.clone();
67 }
68 let hash = read_last_hash_from_file(path);
69 *guard = Some(hash.clone());
70 hash
71}
72
73fn read_last_hash_from_file(path: &PathBuf) -> String {
74 let Ok(file) = fs::File::open(path) else {
75 return "genesis".to_string();
76 };
77 let reader = std::io::BufReader::new(file);
78 let mut last_hash = "genesis".to_string();
79 for line in reader.lines().map_while(Result::ok) {
80 if let Ok(entry) = serde_json::from_str::<AuditEntry>(&line) {
81 last_hash = entry.entry_hash;
82 }
83 }
84 last_hash
85}
86
87fn compute_entry_hash(prev_hash: &str, data_json: &str) -> String {
88 let mut hasher = Sha256::new();
89 hasher.update(prev_hash.as_bytes());
90 hasher.update(data_json.as_bytes());
91 format!("{:x}", hasher.finalize())
92}
93
94pub fn record(data: AuditEntryData) {
95 let Some(path) = trail_path() else { return };
96 let prev_hash = init_last_hash(&path);
97
98 let partial = serde_json::json!({
99 "agent_id": data.agent_id,
100 "tool": data.tool,
101 "action": data.action,
102 "input_hash": data.input_hash,
103 "output_tokens": data.output_tokens,
104 "role": data.role,
105 "event_type": data.event_type,
106 });
107 let data_json = serde_json::to_string(&partial).unwrap_or_default();
108 let entry_hash = compute_entry_hash(&prev_hash, &data_json);
109
110 let entry = AuditEntry {
111 timestamp: chrono::Utc::now().to_rfc3339(),
112 agent_id: data.agent_id,
113 tool: data.tool,
114 action: data.action,
115 input_hash: data.input_hash,
116 output_tokens: data.output_tokens,
117 role: data.role,
118 event_type: data.event_type,
119 prev_hash,
120 entry_hash: entry_hash.clone(),
121 };
122
123 if let Ok(line) = serde_json::to_string(&entry) {
124 if let Ok(mut file) = OpenOptions::new().create(true).append(true).open(&path) {
125 let _ = writeln!(file, "{line}");
126 }
127 }
128
129 if let Ok(mut guard) = LAST_HASH.lock() {
130 *guard = Some(entry_hash);
131 }
132}
133
134pub fn load_recent(limit: usize) -> Vec<AuditEntry> {
135 let Some(path) = trail_path() else {
136 return Vec::new();
137 };
138 let Ok(file) = fs::File::open(&path) else {
139 return Vec::new();
140 };
141 let reader = std::io::BufReader::new(file);
142 let entries: Vec<AuditEntry> = reader
143 .lines()
144 .map_while(Result::ok)
145 .filter_map(|line| serde_json::from_str(&line).ok())
146 .collect();
147 let skip = entries.len().saturating_sub(limit);
148 entries.into_iter().skip(skip).collect()
149}
150
151pub fn verify_chain() -> ChainVerifyResult {
152 let Some(path) = trail_path() else {
153 return ChainVerifyResult {
154 total_entries: 0,
155 valid: true,
156 first_invalid_at: None,
157 };
158 };
159 let Ok(file) = fs::File::open(&path) else {
160 return ChainVerifyResult {
161 total_entries: 0,
162 valid: true,
163 first_invalid_at: None,
164 };
165 };
166 let reader = std::io::BufReader::new(file);
167 let mut prev_hash = "genesis".to_string();
168 let mut total = 0usize;
169
170 for line in reader.lines().map_while(Result::ok) {
171 let entry: AuditEntry = match serde_json::from_str(&line) {
172 Ok(e) => e,
173 Err(_) => {
174 return ChainVerifyResult {
175 total_entries: total,
176 valid: false,
177 first_invalid_at: Some(total),
178 }
179 }
180 };
181
182 if entry.prev_hash != prev_hash {
183 return ChainVerifyResult {
184 total_entries: total,
185 valid: false,
186 first_invalid_at: Some(total),
187 };
188 }
189
190 let partial = serde_json::json!({
191 "agent_id": entry.agent_id,
192 "tool": entry.tool,
193 "action": entry.action,
194 "input_hash": entry.input_hash,
195 "output_tokens": entry.output_tokens,
196 "role": entry.role,
197 "event_type": entry.event_type,
198 });
199 let data_json = serde_json::to_string(&partial).unwrap_or_default();
200 let expected = compute_entry_hash(&prev_hash, &data_json);
201
202 if entry.entry_hash != expected {
203 return ChainVerifyResult {
204 total_entries: total,
205 valid: false,
206 first_invalid_at: Some(total),
207 };
208 }
209
210 prev_hash = entry.entry_hash;
211 total += 1;
212 }
213
214 ChainVerifyResult {
215 total_entries: total,
216 valid: true,
217 first_invalid_at: None,
218 }
219}
220
221pub fn hash_input(args: &serde_json::Map<String, serde_json::Value>) -> String {
222 let serialized = serde_json::to_string(args).unwrap_or_default();
223 let mut hasher = Sha256::new();
224 hasher.update(serialized.as_bytes());
225 format!("{:x}", hasher.finalize())
226}