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