1use crate::query::TRonQuery;
10use bote::dispatch::ToolHandler;
11use bote::registry::{ToolDef, ToolSchema};
12use std::collections::HashMap;
13use std::sync::Arc;
14
15pub 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
77pub fn status_handler(query: TRonQuery) -> ToolHandler {
79 Arc::new(move |_params| {
80 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
101pub 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
132pub 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
161pub 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}