Skip to main content

t_ron/
tools.rs

1//! MCP tools — t-ron's own tools registered with bote for security queries.
2//!
3//! These tools are designed for the T.Ron personality in SecureYeoman:
4//! - `tron_status` — overall security status
5//! - `tron_risk` — per-agent risk score
6//! - `tron_audit` — recent audit events
7//! - `tron_policy` — current policy summary
8
9use crate::query::TRonQuery;
10use bote::dispatch::ToolHandler;
11use bote::registry::{ToolDef, ToolSchema};
12use std::collections::HashMap;
13use std::sync::Arc;
14
15/// Tool definitions for t-ron's MCP tools.
16pub fn tool_defs() -> Vec<ToolDef> {
17    vec![
18        ToolDef {
19            name: "tron_status".into(),
20            description: "Get overall security status: total events, denials, and system health."
21                .into(),
22            input_schema: ToolSchema {
23                schema_type: "object".into(),
24                properties: HashMap::new(),
25                required: vec![],
26            },
27        },
28        ToolDef {
29            name: "tron_risk".into(),
30            description:
31                "Get risk score for an agent (0.0 = trusted, 1.0 = hostile). Requires agent_id."
32                    .into(),
33            input_schema: ToolSchema {
34                schema_type: "object".into(),
35                properties: HashMap::from([(
36                    "agent_id".into(),
37                    serde_json::json!({"type": "string", "description": "Agent to score"}),
38                )]),
39                required: vec!["agent_id".into()],
40            },
41        },
42        ToolDef {
43            name: "tron_audit".into(),
44            description:
45                "Get recent security events. Optional agent_id filter and limit (default 20)."
46                    .into(),
47            input_schema: ToolSchema {
48                schema_type: "object".into(),
49                properties: HashMap::from([
50                    (
51                        "agent_id".into(),
52                        serde_json::json!({"type": "string", "description": "Filter by agent"}),
53                    ),
54                    (
55                        "limit".into(),
56                        serde_json::json!({"type": "integer", "description": "Max events to return", "default": 20}),
57                    ),
58                ]),
59                required: vec![],
60            },
61        },
62        ToolDef {
63            name: "tron_policy".into(),
64            description: "Load or reload policy from a TOML string.".into(),
65            input_schema: ToolSchema {
66                schema_type: "object".into(),
67                properties: HashMap::from([(
68                    "toml".into(),
69                    serde_json::json!({"type": "string", "description": "Policy TOML content"}),
70                )]),
71                required: vec!["toml".into()],
72            },
73        },
74    ]
75}
76
77/// Create handler for `tron_status`.
78pub fn status_handler(query: TRonQuery) -> ToolHandler {
79    Arc::new(move |_params| {
80        // block_on since bote handlers are sync; TRonQuery is Send+Sync (Arc internals)
81        // so no mutex needed — &self methods only.
82        let rt = tokio::runtime::Handle::current();
83        rt.block_on(async {
84            let total = query.total_events().await;
85            let denials = query.total_denials().await;
86            serde_json::json!({
87                "content": [{
88                    "type": "text",
89                    "text": serde_json::to_string_pretty(&serde_json::json!({
90                        "total_events": total,
91                        "total_denials": denials,
92                        "denial_rate": if total > 0 { denials as f64 / total as f64 } else { 0.0 },
93                        "status": if denials == 0 { "clean" } else { "active" }
94                    })).unwrap()
95                }]
96            })
97        })
98    })
99}
100
101/// Create handler for `tron_risk`.
102pub fn risk_handler(query: TRonQuery) -> ToolHandler {
103    Arc::new(move |params| {
104        let agent_id = params
105            .get("agent_id")
106            .and_then(|v| v.as_str())
107            .unwrap_or("")
108            .to_string();
109        let rt = tokio::runtime::Handle::current();
110        rt.block_on(async {
111            let score = query.agent_risk_score(&agent_id).await;
112            let level = match score {
113                s if s >= 0.8 => "critical",
114                s if s >= 0.5 => "high",
115                s if s >= 0.2 => "medium",
116                _ => "low",
117            };
118            serde_json::json!({
119                "content": [{
120                    "type": "text",
121                    "text": serde_json::to_string_pretty(&serde_json::json!({
122                        "agent_id": agent_id,
123                        "risk_score": score,
124                        "risk_level": level
125                    })).unwrap()
126                }]
127            })
128        })
129    })
130}
131
132/// Create handler for `tron_audit`.
133pub fn audit_handler(query: TRonQuery) -> ToolHandler {
134    Arc::new(move |params| {
135        let agent_id = params
136            .get("agent_id")
137            .and_then(|v| v.as_str())
138            .map(|s| s.to_string());
139        let limit = params
140            .get("limit")
141            .and_then(|v| v.as_u64())
142            .unwrap_or(20)
143            .min(1000) as usize;
144        let rt = tokio::runtime::Handle::current();
145        rt.block_on(async {
146            let events = if let Some(ref aid) = agent_id {
147                query.agent_audit(aid, limit).await
148            } else {
149                query.recent_events(limit).await
150            };
151            serde_json::json!({
152                "content": [{
153                    "type": "text",
154                    "text": serde_json::to_string_pretty(&events).unwrap()
155                }]
156            })
157        })
158    })
159}
160
161/// Create handler for `tron_policy`.
162pub fn policy_handler(tron: &crate::TRon) -> ToolHandler {
163    let policy = tron.policy_arc();
164    Arc::new(move |params| {
165        let toml_str = params.get("toml").and_then(|v| v.as_str()).unwrap_or("");
166        if toml_str.trim().is_empty() {
167            return serde_json::json!({
168                "content": [{
169                    "type": "text",
170                    "text": "policy error: empty TOML input"
171                }],
172                "isError": true
173            });
174        }
175        match policy.load_toml(toml_str) {
176            Ok(()) => serde_json::json!({
177                "content": [{
178                    "type": "text",
179                    "text": "policy reloaded successfully"
180                }]
181            }),
182            Err(e) => serde_json::json!({
183                "content": [{
184                    "type": "text",
185                    "text": format!("policy error: {e}")
186                }],
187                "isError": true
188            }),
189        }
190    })
191}
192
193#[cfg(test)]
194mod tests {
195    use super::*;
196
197    #[test]
198    fn tool_defs_all_present() {
199        let defs = tool_defs();
200        assert_eq!(defs.len(), 4);
201        let names: Vec<&str> = defs.iter().map(|d| d.name.as_str()).collect();
202        assert!(names.contains(&"tron_status"));
203        assert!(names.contains(&"tron_risk"));
204        assert!(names.contains(&"tron_audit"));
205        assert!(names.contains(&"tron_policy"));
206    }
207
208    #[test]
209    fn tool_defs_schemas_valid() {
210        for def in tool_defs() {
211            assert_eq!(def.input_schema.schema_type, "object");
212            assert!(!def.description.is_empty());
213        }
214    }
215}