telos_agent/integrations/plugin/
policy_loader.rs1use std::collections::HashMap;
4use std::path::{Path, PathBuf};
5
6use async_trait::async_trait;
7use tokio::io::AsyncWriteExt;
8use tokio::process::Command;
9
10use crate::agent::policies::{Policy, PolicyContext, PolicyOutcome};
11use crate::error::AgentError;
12
13pub struct CommandPolicy {
14 name: String,
15 command: String,
16 args: Vec<String>,
17 timeout_ms: u64,
18 plugin_root: PathBuf,
19 env: HashMap<String, String>,
20}
21
22impl CommandPolicy {
23 pub fn new(
24 name: String,
25 command: String,
26 args: Vec<String>,
27 timeout_ms: u64,
28 plugin_root: PathBuf,
29 env: HashMap<String, String>,
30 ) -> Self {
31 Self { name, command, args, timeout_ms, plugin_root, env }
32 }
33
34 fn resolve(&self, value: &str) -> String {
35 value.replace("${PLUGIN_ROOT}", &self.plugin_root.to_string_lossy())
36 }
37
38 fn command_path(&self) -> String {
39 let command = self.resolve(&self.command);
40 let path = Path::new(&command);
41 if path.is_relative()
42 && (command.starts_with('.') || command.contains('/') || command.contains('\\'))
43 {
44 self.plugin_root.join(path).to_string_lossy().into_owned()
45 } else {
46 command
47 }
48 }
49}
50
51#[async_trait]
52impl Policy for CommandPolicy {
53 fn name(&self) -> &str {
54 &self.name
55 }
56
57 async fn evaluate(&self, context: &PolicyContext) -> Result<PolicyOutcome, AgentError> {
58 let input = serde_json::to_vec(context).map_err(|error| AgentError::ToolExecution {
59 tool: self.name.clone(),
60 message: format!("policy serialization error: {error}"),
61 })?;
62 let mut command = Command::new(self.command_path());
63 command
64 .args(self.args.iter().map(|arg| self.resolve(arg)))
65 .current_dir(&self.plugin_root)
66 .env_clear()
67 .envs(&self.env)
68 .env("PLUGIN_ROOT", &self.plugin_root)
69 .stdin(std::process::Stdio::piped())
70 .stdout(std::process::Stdio::piped())
71 .stderr(std::process::Stdio::piped())
72 .kill_on_drop(true);
73 let mut child = command.spawn().map_err(|error| AgentError::ToolExecution {
74 tool: self.name.clone(),
75 message: format!("policy failed to spawn: {error}"),
76 })?;
77 if let Some(mut stdin) = child.stdin.take() {
78 stdin.write_all(&input).await.map_err(|error| AgentError::ToolExecution {
79 tool: self.name.clone(),
80 message: format!("policy stdin failed: {error}"),
81 })?;
82 }
83 let output = tokio::time::timeout(
84 std::time::Duration::from_millis(self.timeout_ms),
85 child.wait_with_output(),
86 )
87 .await
88 .map_err(|_| AgentError::ToolExecution {
89 tool: self.name.clone(),
90 message: "policy timed out".into(),
91 })?
92 .map_err(|error| AgentError::ToolExecution {
93 tool: self.name.clone(),
94 message: format!("policy failed: {error}"),
95 })?;
96 if !output.status.success() {
97 return Err(AgentError::ToolExecution {
98 tool: self.name.clone(),
99 message: String::from_utf8_lossy(&output.stderr).trim().to_string(),
100 });
101 }
102 serde_json::from_slice(&output.stdout).map_err(|error| AgentError::ToolExecution {
103 tool: self.name.clone(),
104 message: format!("invalid policy outcome: {error}"),
105 })
106 }
107}
108
109#[cfg(test)]
110mod tests {
111 use super::*;
112
113 #[test]
114 fn resolves_plugin_root_in_commands_and_arguments() {
115 let root = PathBuf::from("example-plugin");
116 let policy = CommandPolicy::new(
117 "p".into(),
118 "./bin/check".into(),
119 vec!["${PLUGIN_ROOT}/config.json".into()],
120 1000,
121 root.clone(),
122 HashMap::new(),
123 );
124 assert_eq!(PathBuf::from(policy.command_path()), root.join("./bin/check"));
125 assert_eq!(PathBuf::from(policy.resolve(&policy.args[0])), root.join("config.json"));
126 }
127}