kindly_guard_server/protocol/
claude_code.rs

1// Copyright 2025 Kindly Software Inc.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14//! Claude Code specific MCP protocol extensions
15//!
16//! This module defines the MCP protocol extensions for Claude Code integration,
17//! providing real-time shield status notifications and control methods.
18
19use serde::{Deserialize, Serialize};
20use serde_json::Value;
21
22/// Shield status notification for Claude Code
23#[derive(Debug, Clone, Serialize, Deserialize)]
24#[serde(rename_all = "camelCase")]
25pub struct ShieldStatusNotification {
26    /// Always "2.0" for JSON-RPC
27    pub jsonrpc: String,
28    /// Method name: "shield/status"
29    pub method: String,
30    /// Shield status parameters
31    pub params: ShieldStatusParams,
32}
33
34/// Shield status parameters
35#[derive(Debug, Clone, Serialize, Deserialize)]
36#[serde(rename_all = "camelCase")]
37pub struct ShieldStatusParams {
38    /// Is the shield currently active
39    pub active: bool,
40    /// Is enhanced mode enabled
41    pub enhanced: bool,
42    /// Total number of threats blocked
43    pub threats: u64,
44    /// Threat detection rate per minute
45    pub threat_rate: f64,
46    /// Last detected threat information
47    #[serde(skip_serializing_if = "Option::is_none")]
48    pub last_threat: Option<LastThreatInfo>,
49    /// Performance metrics
50    pub performance: PerformanceMetrics,
51}
52
53/// Information about the last detected threat
54#[derive(Debug, Clone, Serialize, Deserialize)]
55#[serde(rename_all = "camelCase")]
56pub struct LastThreatInfo {
57    /// Type of threat detected
58    #[serde(rename = "type")]
59    pub threat_type: String,
60    /// Severity level
61    pub severity: ThreatSeverity,
62    /// Unix timestamp in milliseconds
63    pub timestamp: u64,
64    /// Human-readable description
65    pub description: String,
66}
67
68/// Threat severity levels
69#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
70#[serde(rename_all = "lowercase")]
71pub enum ThreatSeverity {
72    Low,
73    Medium,
74    High,
75    Critical,
76}
77
78/// Performance metrics for Claude Code monitoring
79#[derive(Debug, Clone, Serialize, Deserialize)]
80#[serde(rename_all = "camelCase")]
81pub struct PerformanceMetrics {
82    /// Last scan time in microseconds
83    pub scan_time_us: u64,
84    /// Current queue depth
85    pub queue_depth: usize,
86    /// Memory usage in MB
87    pub memory_mb: f64,
88}
89
90/// Shield control request
91#[derive(Debug, Clone, Serialize, Deserialize)]
92#[serde(rename_all = "camelCase")]
93pub struct ShieldControlRequest {
94    /// Control action to perform
95    pub action: ShieldControlAction,
96    /// Optional duration in milliseconds
97    #[serde(skip_serializing_if = "Option::is_none")]
98    pub duration: Option<u64>,
99}
100
101/// Shield control actions
102#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
103#[serde(rename_all = "lowercase")]
104pub enum ShieldControlAction {
105    /// Temporarily pause shield
106    Pause,
107    /// Resume shield operation
108    Resume,
109    /// Reset statistics
110    Reset,
111    /// Enable enhanced mode
112    Enhance,
113}
114
115/// Shield control response
116#[derive(Debug, Clone, Serialize, Deserialize)]
117#[serde(rename_all = "camelCase")]
118pub struct ShieldControlResponse {
119    /// Whether the action was successful
120    pub success: bool,
121    /// New shield state after action
122    pub state: ShieldState,
123    /// Optional message
124    #[serde(skip_serializing_if = "Option::is_none")]
125    pub message: Option<String>,
126}
127
128/// Current shield state
129#[derive(Debug, Clone, Serialize, Deserialize)]
130#[serde(rename_all = "camelCase")]
131pub struct ShieldState {
132    /// Is shield active
133    pub active: bool,
134    /// Is paused
135    pub paused: bool,
136    /// Enhanced mode enabled
137    pub enhanced: bool,
138    /// Pause end time (if paused)
139    #[serde(skip_serializing_if = "Option::is_none")]
140    pub pause_until: Option<u64>,
141}
142
143/// Shield info request parameters
144#[derive(Debug, Clone, Serialize, Deserialize)]
145#[serde(rename_all = "camelCase")]
146pub struct ShieldInfoParams {
147    /// Request detailed information
148    #[serde(default)]
149    pub detailed: bool,
150}
151
152/// Shield info response
153#[derive(Debug, Clone, Serialize, Deserialize)]
154#[serde(rename_all = "camelCase")]
155pub struct ShieldInfoResponse {
156    /// Shield version
157    pub version: String,
158    /// Current state
159    pub state: ShieldState,
160    /// Statistics
161    pub stats: ShieldStatistics,
162    /// Configuration (if detailed requested)
163    #[serde(skip_serializing_if = "Option::is_none")]
164    pub config: Option<ShieldConfig>,
165    /// Threat patterns (if detailed requested)
166    #[serde(skip_serializing_if = "Option::is_none")]
167    pub patterns: Option<Vec<ThreatPattern>>,
168}
169
170/// Shield statistics
171#[derive(Debug, Clone, Serialize, Deserialize)]
172#[serde(rename_all = "camelCase")]
173pub struct ShieldStatistics {
174    /// Total threats blocked
175    pub threats_blocked: u64,
176    /// Threats by type
177    pub threats_by_type: std::collections::HashMap<String, u64>,
178    /// Total scans performed
179    pub total_scans: u64,
180    /// Average scan time in microseconds
181    pub avg_scan_time_us: u64,
182    /// Uptime in seconds
183    pub uptime_seconds: u64,
184    /// Memory usage in MB
185    pub memory_usage_mb: f64,
186}
187
188/// Shield configuration info
189#[derive(Debug, Clone, Serialize, Deserialize)]
190#[serde(rename_all = "camelCase")]
191pub struct ShieldConfig {
192    /// Scanner sensitivity level
193    pub sensitivity: String,
194    /// Enabled threat detectors
195    pub enabled_detectors: Vec<String>,
196    /// Rate limiting enabled
197    pub rate_limiting: bool,
198    /// Max threats per minute before blocking
199    pub max_threat_rate: u64,
200}
201
202/// Threat pattern information
203#[derive(Debug, Clone, Serialize, Deserialize)]
204#[serde(rename_all = "camelCase")]
205pub struct ThreatPattern {
206    /// Pattern name
207    pub name: String,
208    /// Pattern type
209    #[serde(rename = "type")]
210    pub pattern_type: String,
211    /// Is enabled
212    pub enabled: bool,
213    /// Detection count
214    pub detections: u64,
215}
216
217/// Claude Code specific errors
218#[derive(Debug, Clone, Serialize, Deserialize)]
219#[serde(rename_all = "camelCase")]
220pub struct ClaudeCodeError {
221    /// Error code
222    pub code: ClaudeCodeErrorCode,
223    /// Error message
224    pub message: String,
225    /// Additional details
226    #[serde(skip_serializing_if = "Option::is_none")]
227    pub details: Option<Value>,
228}
229
230/// Claude Code error codes
231#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
232pub enum ClaudeCodeErrorCode {
233    /// Shield is not available
234    ShieldUnavailable = -40001,
235    /// Invalid control action
236    InvalidAction = -40002,
237    /// Operation timeout
238    OperationTimeout = -40003,
239    /// Enhanced mode not available
240    EnhancedModeUnavailable = -40004,
241    /// Rate limit exceeded
242    RateLimitExceeded = -40005,
243}
244
245/// Binary protocol message header (enhanced mode)
246#[cfg(feature = "enhanced")]
247#[repr(C, packed)]
248pub struct BinaryMessageHeader {
249    /// Magic number: 0x4B475344 ('KGSD')
250    pub magic: u32,
251    /// Protocol version
252    pub version: u16,
253    /// Message type
254    pub msg_type: u16,
255    /// Payload length
256    pub payload_len: u32,
257    /// Timestamp (nanoseconds since epoch)
258    pub timestamp: u64,
259    /// Sequence number
260    pub sequence: u32,
261    /// Checksum
262    pub checksum: u32,
263}
264
265#[cfg(feature = "enhanced")]
266impl BinaryMessageHeader {
267    pub const MAGIC: u32 = 0x4B475344; // 'KGSD'
268    pub const VERSION: u16 = 1;
269
270    pub const MSG_TYPE_STATUS: u16 = 1;
271    pub const MSG_TYPE_THREAT: u16 = 2;
272    pub const MSG_TYPE_CONTROL: u16 = 3;
273    pub const MSG_TYPE_PERF: u16 = 4;
274}
275
276/// Helper to create shield status notification
277pub fn create_status_notification(params: ShieldStatusParams) -> ShieldStatusNotification {
278    ShieldStatusNotification {
279        jsonrpc: "2.0".to_string(),
280        method: "shield/status".to_string(),
281        params,
282    }
283}
284
285/// Helper to convert threat to severity
286pub fn threat_to_severity(threat: &crate::scanner::Threat) -> ThreatSeverity {
287    match &threat.threat_type {
288        crate::scanner::ThreatType::UnicodeInvisible => ThreatSeverity::Medium,
289        crate::scanner::ThreatType::UnicodeBiDi => ThreatSeverity::High,
290        crate::scanner::ThreatType::UnicodeHomograph => ThreatSeverity::Medium,
291        crate::scanner::ThreatType::UnicodeControl => ThreatSeverity::Medium,
292        crate::scanner::ThreatType::PromptInjection => ThreatSeverity::High,
293        crate::scanner::ThreatType::CommandInjection => ThreatSeverity::Critical,
294        crate::scanner::ThreatType::PathTraversal => ThreatSeverity::High,
295        crate::scanner::ThreatType::SqlInjection => ThreatSeverity::Critical,
296        crate::scanner::ThreatType::CrossSiteScripting => ThreatSeverity::High,
297        crate::scanner::ThreatType::LdapInjection => ThreatSeverity::High,
298        crate::scanner::ThreatType::XmlInjection => ThreatSeverity::High,
299        crate::scanner::ThreatType::NoSqlInjection => ThreatSeverity::High,
300        crate::scanner::ThreatType::SessionIdExposure => ThreatSeverity::Critical,
301        crate::scanner::ThreatType::ToolPoisoning => ThreatSeverity::Critical,
302        crate::scanner::ThreatType::TokenTheft => ThreatSeverity::Critical,
303        crate::scanner::ThreatType::DosPotential => ThreatSeverity::High,
304        crate::scanner::ThreatType::Custom(_) => ThreatSeverity::Medium,
305    }
306}
307
308#[cfg(test)]
309mod tests {
310    use super::*;
311
312    #[test]
313    fn test_status_notification_serialization() {
314        let notification = create_status_notification(ShieldStatusParams {
315            active: true,
316            enhanced: false,
317            threats: 42,
318            threat_rate: 2.5,
319            last_threat: Some(LastThreatInfo {
320                threat_type: "unicode_bidi".to_string(),
321                severity: ThreatSeverity::High,
322                timestamp: 1234567890,
323                description: "Right-to-left override detected".to_string(),
324            }),
325            performance: PerformanceMetrics {
326                scan_time_us: 123,
327                queue_depth: 5,
328                memory_mb: 45.6,
329            },
330        });
331
332        let json = serde_json::to_string_pretty(&notification).unwrap();
333        assert!(json.contains("\"method\": \"shield/status\""));
334        assert!(json.contains("\"threats\": 42"));
335        assert!(json.contains("\"severity\": \"high\""));
336    }
337
338    #[test]
339    fn test_control_request_deserialization() {
340        let json = r#"{
341            "action": "pause",
342            "duration": 5000
343        }"#;
344
345        let request: ShieldControlRequest = serde_json::from_str(json).unwrap();
346        assert_eq!(request.action, ShieldControlAction::Pause);
347        assert_eq!(request.duration, Some(5000));
348    }
349
350    #[cfg(feature = "enhanced")]
351    #[test]
352    fn test_binary_header() {
353        use std::mem;
354
355        assert_eq!(mem::size_of::<BinaryMessageHeader>(), 28);
356        assert_eq!(BinaryMessageHeader::MAGIC, 0x4B475344);
357    }
358}
kindly_guard_server/protocol/claude_code.rs

kindly_guard_server/protocol/
claude_code.rs
