Scaffoled Script Support

hermes/src/main.rs

@@ -113,7 +113,7 @@ async fn eval(state: Arc<RwLock<State>>) {
         AgentInstructionBody::Script(script) => {
             task::spawn_blocking(async move || {
                 let mut engine = ScriptingEngine::new();
-                engine.execute(&script.source).await;
+                engine.execute(&script.body).await;
             })
             .await;
 

talaria/Cargo.toml

@@ -16,3 +16,4 @@ strum = "0.27"
 strum_macros = "0.27"
 bytesize = "2.0.1"
 rhai = { version = "1.21", featues = ["sync"] }
+toml = "0.8.22"

talaria/src/console.rs

@@ -303,6 +303,15 @@ pub enum GroupCommand {
     None,
 }
 
+#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, EnumProperty, EnumIter)]
+pub enum Param {
+    String(String),
+    Int(i64),
+    Float(f64),
+    Bool(bool),
+    Array(Vec<Param>),
+}
+
 #[derive(Clone, Debug, Serialize, Deserialize, PartialEq, EnumProperty, Default, EnumIter)]
 pub enum RunCommand {
     #[strum(props(
@@ -314,6 +323,7 @@ pub enum RunCommand {
     Script {
         target: Option<TargetIdentifier>,
         script_name: String,
+        params: Vec<Param>,
     },
 
     #[strum(props(
@@ -391,13 +401,12 @@ pub enum CommandError {
     ExpectedRhai,
     #[error("expected shell script")]
     ExpectedShell,
-}
-
-pub enum Token {
-    CommandName { command_name: String },
-    AgentID { id: u64 },
-    AgentNickname { nickname: String },
-    GroupIdentifier { identifier: String },
+    #[error("expected script param")]
+    ExpectedScriptParam,
+    #[error("expected closing bracket")]
+    ExpectedClosingBracket,
+    #[error("invalid number")]
+    InvalidNumber,
 }
 
 pub struct Parser {
@@ -477,6 +486,12 @@ impl Parser {
                         tokens.push(current_token.iter().collect());
                         current_token.clear();
                     }
+                } else if char == '[' || char == ']' {
+                    if current_token.len() == 0 {
+                        tokens.push(char.to_string());
+                    } else {
+                        current_token.push(char);
+                    }
                 } else {
                     current_token.push(char);
                 }
@@ -617,7 +632,7 @@ impl Parser {
     ) -> Result<Option<AgentIdentifier>, CommandError> {
         match last_arg {
             true =>
-            // his is a little bit of a hack,
+            // this is a little bit of a hack,
             // but it's very ergonomic to use-- and i'm not sure how to
             // implement this in a nicer way
             //
@@ -863,6 +878,7 @@ impl Parser {
             RunCommand::Script { .. } => RunCommand::Script {
                 target: self.parse_opt_target_ident(false)?,
                 script_name: self.parse_script_name()?,
+                params: self.parse_script_params()?,
             },
             RunCommand::Rhai { .. } => RunCommand::Rhai {
                 target: self.parse_opt_target_ident(false)?,
@@ -885,6 +901,66 @@ impl Parser {
             _ => Err(CommandError::InvalidScriptName),
         }
     }
+
+    pub fn parse_script_params(&mut self) -> Result<Vec<Param>, CommandError> {
+        // FIXME:
+        Ok(vec![])
+        // if self.is_at_end() {
+        //     return Ok(vec![]);
+        // }
+        //
+        // let mut params = vec![];
+        //
+        // loop {
+        //     let token = self.peek(CommandError::ExpectedScriptParam);
+        //     let token = token?;
+        //
+        //     let next_char = token
+        //         .chars()
+        //         .next()
+        //         .ok_or(CommandError::ExpectedClosingBracket)?;
+        //
+        //     match next_char {
+        //         'a'..='z' | 'A'..='Z' => {
+        //             if token == "true" {
+        //                 params.push(Param::Bool(true));
+        //             }
+        //
+        //             if token == "false" {
+        //                 params.push(Param::Bool(false));
+        //             }
+        //
+        //             params.push(self.parse_script_string()?);
+        //         }
+        //         '0'..'9' | '.' | '-' => params.push(self.parse_script_number()?),
+        //         '[' => {}
+        //         ']' => {}
+        //         _ => {}
+        //     }
+        // }
+    }
+
+    pub fn parse_script_number(&mut self) -> Result<Param, CommandError> {
+        let token = self.consume(CommandError::ExpectedScriptParam)?;
+
+        match token.parse::<f64>() {
+            Ok(float) => return Ok(Param::Float(float)),
+            Err(_) => {}
+        }
+
+        match token.parse::<i64>() {
+            Ok(int) => return Ok(Param::Int(int)),
+            Err(_) => {}
+        }
+
+        return Err(CommandError::InvalidNumber);
+    }
+
+    pub fn parse_script_string(&mut self) -> Result<Param, CommandError> {
+        let token = self.consume(CommandError::ExpectedScriptParam);
+
+        Ok(Param::String(token?.to_string()))
+    }
 }
 
 pub struct Console {

talaria/src/protocol.rs

@@ -1,23 +1,9 @@
+use crate::scripting::Script;
+
 use anyhow::Result;
 use bincode::{Decode, Encode};
 use serde::{Deserialize, Serialize};
 
-#[derive(Clone, Debug, Encode, Decode, Serialize, Deserialize)]
-pub struct Script {
-    pub source: String,
-    pub description: String,
-    pub title: String,
-}
-
-impl ToString for Script {
-    fn to_string(&self) -> String {
-        format!(
-            "Title: {}\nDescription: {}\n\n{}",
-            self.title, self.description, self.source
-        )
-    }
-}
-
 #[derive(Encode, Decode, Serialize, Deserialize, Clone, Debug, PartialEq)]
 pub struct OS {
     pub os_type: OSType,
@@ -136,7 +122,7 @@ impl AgentInstructionBody {
                 format!("Command: {}\nArgs: {:#?}", command, args)
             }
             AgentInstructionBody::Ok => String::from("None"),
-            AgentInstructionBody::Script(script) => script.title.clone(),
+            AgentInstructionBody::Script(script) => script.name.clone(),
             AgentInstructionBody::Rhai(source) => source.clone(),
             AgentInstructionBody::Kill => String::from("None"),
         }

talaria/src/scripting.rs

@@ -1,5 +1,8 @@
-use anyhow::Result;
+use anyhow::{anyhow, Result};
+use bincode::{Decode, Encode};
+use bytesize::Display;
 use rhai::{plugin::*, Scope};
+use serde::{Deserialize, Serialize};
 
 use crate::stdlib::*;
 
@@ -44,3 +47,119 @@ impl ScriptingEngine {
         &self.engine
     }
 }
+
+#[derive(Clone, Deserialize, Serialize, Encode, Decode, Debug)]
+pub struct Script {
+    pub name: String,
+    pub description: String,
+    pub params: Vec<Param>,
+    #[serde(skip)]
+    pub body: String,
+}
+
+impl Script {
+    pub fn signature(&self) -> String {
+        let parameters = self
+            .params
+            .clone()
+            .into_iter()
+            .map(|p| format!("[{}: {}]", p.name, p.r#type.to_string()))
+            .collect::<Vec<String>>()
+            .join(" ");
+
+        let example = format!(
+            "{} {}",
+            self.name,
+            self.params
+                .clone()
+                .into_iter()
+                .map(|p| format!("{}", p.placeholder))
+                .collect::<Vec<String>>()
+                .join(" ")
+        );
+
+        format!(
+            "Name: {}\nDescription: {}\nParameters: {}\nExample: {}",
+            self.name, self.description, parameters, example
+        )
+    }
+
+    pub fn from_str(input: &str) -> Result<Script> {
+        let input = input.trim();
+
+        let lines = input.split('\n').collect::<Vec<&str>>();
+
+        let line = match lines.get(0) {
+            Some(line) => line,
+            None => return Err(anyhow!("Script cannot be empty")),
+        };
+
+        if *line != "---" {
+            return Err(anyhow!("Script must start with metadata"));
+        }
+
+        let mut index = 1;
+
+        let header_end = loop {
+            let line = match lines.get(index) {
+                Some(line) => line,
+                None => return Err(anyhow!("Script metadata is never closed")),
+            };
+
+            if *line == "---" {
+                break index;
+            }
+
+            index += 1;
+        };
+
+        if header_end == 1 {
+            return Err(anyhow!("Script metadata cannot be empty"));
+        }
+
+        if lines.len() < header_end + 1 {
+            return Err(anyhow!("Script must contain body"));
+        }
+
+        let header = lines[1..header_end].join("\n");
+
+        let mut script = match toml::from_str::<Script>(&header) {
+            Ok(script) => script,
+            Err(err) => return Err(anyhow!(err.to_string())),
+        };
+
+        script.body = lines[header_end + 1..].join("\n");
+
+        Ok(script)
+    }
+}
+
+#[derive(Clone, Debug, Deserialize, Serialize, Encode, Decode)]
+pub struct Param {
+    pub name: String,
+    pub arg_name: String,
+    pub description: String,
+    pub r#type: ParamType,
+    pub placeholder: String,
+}
+
+#[derive(Clone, Deserialize, Serialize, Encode, Decode, Debug)]
+#[serde(rename_all = "lowercase")]
+pub enum ParamType {
+    String,
+    Int,
+    Float,
+    Bool,
+    Array,
+}
+impl ToString for ParamType {
+    fn to_string(&self) -> String {
+        match self {
+            ParamType::String => String::from("string"),
+            ParamType::Int => String::from("int"),
+            ParamType::Float => String::from("float"),
+            ParamType::Bool => String::from("bool"),
+            ParamType::Array => String::from("array"),
+        }
+    }
+}

tartarus/Cargo.toml

@@ -14,8 +14,8 @@ log = "0.4"
 env_logger = "0.11.8"
 rustyline = { version = "15.0.0", features = ["derive"]}
 anyhow = "1.0"
-toml = "0.8"
 sysinfo = "0.34.2"
 strum = "0.27.1"
 structz-macros = "0.2.0"
 figment = {version = "0.10", features = ["toml", "env"]}
+toml = "0.8.22"