Skip to main content

cc_audit/proxy/
logger.rs

1//! JSONL logger for proxy traffic.
2
3use 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/// Direction of the proxied message.
11#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
12#[serde(rename_all = "lowercase")]
13pub enum MessageDirection {
14    /// Request from client to server
15    Request,
16    /// Response from server to client
17    Response,
18}
19
20/// Log entry for a proxied message.
21#[derive(Debug, Clone, Serialize, Deserialize)]
22pub struct ProxyLog {
23    /// Timestamp of the log entry
24    pub timestamp: String,
25
26    /// Direction of the message
27    pub direction: MessageDirection,
28
29    /// JSON-RPC method (if applicable)
30    #[serde(skip_serializing_if = "Option::is_none")]
31    pub method: Option<String>,
32
33    /// Findings from security analysis
34    #[serde(skip_serializing_if = "Vec::is_empty")]
35    pub findings: Vec<FindingSummary>,
36
37    /// Action taken
38    pub action: String,
39
40    /// Client address
41    #[serde(skip_serializing_if = "Option::is_none")]
42    pub client_addr: Option<String>,
43
44    /// Message size in bytes
45    pub size: usize,
46}
47
48/// Summary of a finding for logging.
49#[derive(Debug, Clone, Serialize, Deserialize)]
50pub struct FindingSummary {
51    /// Rule ID
52    pub id: String,
53    /// Severity level
54    pub severity: String,
55    /// Message
56    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/// Logger for proxy traffic.
70#[derive(Default)]
71pub struct ProxyLogger {
72    /// File writer (if logging to file)
73    writer: Option<Mutex<BufWriter<File>>>,
74
75    /// Verbose mode (log to stderr)
76    verbose: bool,
77}
78
79impl ProxyLogger {
80    /// Create a new logger with optional file output.
81    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    /// Log a message.
93    pub fn log(&self, entry: &ProxyLog) {
94        // Serialize to JSON
95        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        // Write to file if configured
104        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        // Print to stderr in verbose mode
112        if self.verbose {
113            eprintln!("[PROXY] {}", json);
114        }
115    }
116
117    /// Log a request.
118    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    /// Log a response.
140    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        // Read the file
183        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        // Should not panic
215        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        // Read the file
234        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        // Read the file
268        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        // Read the file
283        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        // Read the file
298        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}