1use crate::rules::Finding;
4use serde::{Deserialize, Serialize};
5use std::fs::{File, OpenOptions};
6use std::io::{BufWriter, Write};
7use std::path::Path;
8use std::sync::Mutex;
9
10#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
12#[serde(rename_all = "lowercase")]
13pub enum MessageDirection {
14 Request,
16 Response,
18}
19
20#[derive(Debug, Clone, Serialize, Deserialize)]
22pub struct ProxyLog {
23 pub timestamp: String,
25
26 pub direction: MessageDirection,
28
29 #[serde(skip_serializing_if = "Option::is_none")]
31 pub method: Option<String>,
32
33 #[serde(skip_serializing_if = "Vec::is_empty")]
35 pub findings: Vec<FindingSummary>,
36
37 pub action: String,
39
40 #[serde(skip_serializing_if = "Option::is_none")]
42 pub client_addr: Option<String>,
43
44 pub size: usize,
46}
47
48#[derive(Debug, Clone, Serialize, Deserialize)]
50pub struct FindingSummary {
51 pub id: String,
53 pub severity: String,
55 pub message: String,
57}
58
59impl From<&Finding> for FindingSummary {
60 fn from(f: &Finding) -> Self {
61 Self {
62 id: f.id.clone(),
63 severity: format!("{:?}", f.severity).to_lowercase(),
64 message: f.message.clone(),
65 }
66 }
67}
68
69#[derive(Default)]
71pub struct ProxyLogger {
72 writer: Option<Mutex<BufWriter<File>>>,
74
75 verbose: bool,
77}
78
79impl ProxyLogger {
80 pub fn new(log_path: Option<&Path>, verbose: bool) -> std::io::Result<Self> {
82 let writer = if let Some(path) = log_path {
83 let file = OpenOptions::new().create(true).append(true).open(path)?;
84 Some(Mutex::new(BufWriter::new(file)))
85 } else {
86 None
87 };
88
89 Ok(Self { writer, verbose })
90 }
91
92 pub fn log(&self, entry: &ProxyLog) {
94 let json = match serde_json::to_string(entry) {
96 Ok(j) => j,
97 Err(e) => {
98 eprintln!("Failed to serialize log entry: {}", e);
99 return;
100 }
101 };
102
103 if let Some(ref writer) = self.writer
105 && let Ok(mut w) = writer.lock()
106 {
107 let _ = writeln!(w, "{}", json);
108 let _ = w.flush();
109 }
110
111 if self.verbose {
113 eprintln!("[PROXY] {}", json);
114 }
115 }
116
117 pub fn log_request(
119 &self,
120 method: Option<&str>,
121 findings: &[Finding],
122 action: &str,
123 client_addr: Option<&str>,
124 size: usize,
125 ) {
126 let entry = ProxyLog {
127 timestamp: chrono::Utc::now().to_rfc3339(),
128 direction: MessageDirection::Request,
129 method: method.map(|s| s.to_string()),
130 findings: findings.iter().map(FindingSummary::from).collect(),
131 action: action.to_string(),
132 client_addr: client_addr.map(|s| s.to_string()),
133 size,
134 };
135
136 self.log(&entry);
137 }
138
139 pub fn log_response(
141 &self,
142 method: Option<&str>,
143 findings: &[Finding],
144 action: &str,
145 client_addr: Option<&str>,
146 size: usize,
147 ) {
148 let entry = ProxyLog {
149 timestamp: chrono::Utc::now().to_rfc3339(),
150 direction: MessageDirection::Response,
151 method: method.map(|s| s.to_string()),
152 findings: findings.iter().map(FindingSummary::from).collect(),
153 action: action.to_string(),
154 client_addr: client_addr.map(|s| s.to_string()),
155 size,
156 };
157
158 self.log(&entry);
159 }
160}
161
162#[cfg(test)]
163mod tests {
164 use super::*;
165 use tempfile::TempDir;
166
167 #[test]
168 fn test_log_to_file() {
169 let temp_dir = TempDir::new().unwrap();
170 let log_path = temp_dir.path().join("proxy.jsonl");
171
172 let logger = ProxyLogger::new(Some(&log_path), false).unwrap();
173
174 logger.log_request(
175 Some("tools/call"),
176 &[],
177 "allowed",
178 Some("127.0.0.1:12345"),
179 100,
180 );
181
182 let content = std::fs::read_to_string(&log_path).unwrap();
184 assert!(content.contains("tools/call"));
185 assert!(content.contains("allowed"));
186 assert!(content.contains("request"));
187 }
188
189 #[test]
190 fn test_finding_summary() {
191 use crate::rules::{Category, Severity};
192 use crate::test_utils::fixtures::create_finding;
193
194 let finding = create_finding(
195 "EX-001",
196 Severity::High,
197 Category::Exfiltration,
198 "Test finding",
199 "test.md",
200 1,
201 );
202
203 let summary = FindingSummary::from(&finding);
204
205 assert_eq!(summary.id, "EX-001");
206 assert_eq!(summary.severity, "high");
207 assert!(summary.message.contains("test message"));
208 }
209
210 #[test]
211 fn test_default_logger() {
212 let logger = ProxyLogger::default();
213
214 logger.log_request(None, &[], "allowed", None, 0);
216 }
217
218 #[test]
219 fn test_log_response() {
220 let temp_dir = TempDir::new().unwrap();
221 let log_path = temp_dir.path().join("proxy.jsonl");
222
223 let logger = ProxyLogger::new(Some(&log_path), false).unwrap();
224
225 logger.log_response(
226 Some("tools/call"),
227 &[],
228 "allowed",
229 Some("127.0.0.1:12345"),
230 100,
231 );
232
233 let content = std::fs::read_to_string(&log_path).unwrap();
235 assert!(content.contains("tools/call"));
236 assert!(content.contains("allowed"));
237 assert!(content.contains("response"));
238 }
239
240 #[test]
241 fn test_log_with_findings() {
242 use crate::rules::{Category, Severity};
243 use crate::test_utils::fixtures::create_finding;
244
245 let temp_dir = TempDir::new().unwrap();
246 let log_path = temp_dir.path().join("proxy.jsonl");
247
248 let logger = ProxyLogger::new(Some(&log_path), false).unwrap();
249
250 let finding = create_finding(
251 "EX-001",
252 Severity::High,
253 Category::Exfiltration,
254 "test",
255 "test.md",
256 1,
257 );
258
259 logger.log_request(
260 Some("tools/call"),
261 &[finding],
262 "blocked",
263 Some("127.0.0.1:12345"),
264 100,
265 );
266
267 let content = std::fs::read_to_string(&log_path).unwrap();
269 assert!(content.contains("EX-001"));
270 assert!(content.contains("blocked"));
271 }
272
273 #[test]
274 fn test_log_without_method() {
275 let temp_dir = TempDir::new().unwrap();
276 let log_path = temp_dir.path().join("proxy.jsonl");
277
278 let logger = ProxyLogger::new(Some(&log_path), false).unwrap();
279
280 logger.log_request(None, &[], "allowed", Some("127.0.0.1:12345"), 100);
281
282 let content = std::fs::read_to_string(&log_path).unwrap();
284 assert!(content.contains("request"));
285 assert!(!content.contains("method"));
286 }
287
288 #[test]
289 fn test_log_without_client_addr() {
290 let temp_dir = TempDir::new().unwrap();
291 let log_path = temp_dir.path().join("proxy.jsonl");
292
293 let logger = ProxyLogger::new(Some(&log_path), false).unwrap();
294
295 logger.log_request(Some("test"), &[], "allowed", None, 100);
296
297 let content = std::fs::read_to_string(&log_path).unwrap();
299 assert!(content.contains("test"));
300 assert!(!content.contains("client_addr"));
301 }
302
303 #[test]
304 fn test_message_direction_serialization() {
305 let request_json = serde_json::to_string(&MessageDirection::Request).unwrap();
306 assert_eq!(request_json, "\"request\"");
307
308 let response_json = serde_json::to_string(&MessageDirection::Response).unwrap();
309 assert_eq!(response_json, "\"response\"");
310 }
311
312 #[test]
313 fn test_message_direction_deserialization() {
314 let request: MessageDirection = serde_json::from_str("\"request\"").unwrap();
315 assert_eq!(request, MessageDirection::Request);
316
317 let response: MessageDirection = serde_json::from_str("\"response\"").unwrap();
318 assert_eq!(response, MessageDirection::Response);
319 }
320
321 #[test]
322 fn test_proxy_log_serialization() {
323 let log = ProxyLog {
324 timestamp: "2024-01-01T00:00:00Z".to_string(),
325 direction: MessageDirection::Request,
326 method: Some("test".to_string()),
327 findings: vec![],
328 action: "allowed".to_string(),
329 client_addr: Some("127.0.0.1:8080".to_string()),
330 size: 100,
331 };
332
333 let json = serde_json::to_string(&log).unwrap();
334 assert!(json.contains("2024-01-01"));
335 assert!(json.contains("request"));
336 assert!(json.contains("test"));
337 assert!(json.contains("allowed"));
338 }
339
340 #[test]
341 fn test_proxy_log_deserialization() {
342 let json = r#"{
343 "timestamp": "2024-01-01T00:00:00Z",
344 "direction": "response",
345 "method": "tools/call",
346 "findings": [],
347 "action": "logged",
348 "client_addr": "127.0.0.1:9999",
349 "size": 200
350 }"#;
351
352 let log: ProxyLog = serde_json::from_str(json).unwrap();
353 assert_eq!(log.timestamp, "2024-01-01T00:00:00Z");
354 assert_eq!(log.direction, MessageDirection::Response);
355 assert_eq!(log.method, Some("tools/call".to_string()));
356 assert_eq!(log.action, "logged");
357 assert_eq!(log.size, 200);
358 }
359
360 #[test]
361 fn test_finding_summary_serialization() {
362 let summary = FindingSummary {
363 id: "TEST-001".to_string(),
364 severity: "high".to_string(),
365 message: "Test message".to_string(),
366 };
367
368 let json = serde_json::to_string(&summary).unwrap();
369 assert!(json.contains("TEST-001"));
370 assert!(json.contains("high"));
371 assert!(json.contains("Test message"));
372 }
373}