telos_agent/tools/command_security/powershell/
parser.rs

1//! PowerShell parser wrapper around `tree-sitter-pwsh`.
2
3#[derive(Debug, Clone, PartialEq, Eq)]
4pub struct ParsedPowerShellCommand {
5    commands: Vec<ParsedCommandElement>,
6}
7
8#[derive(Debug, Clone, PartialEq, Eq)]
9pub struct ParsedCommandElement {
10    pub name: String,
11    pub args: Vec<String>,
12    pub dynamic: bool,
13}
14
15impl ParsedPowerShellCommand {
16    pub fn commands(&self) -> &[ParsedCommandElement] {
17        &self.commands
18    }
19}
20
21pub fn parse(command: &str) -> Result<ParsedPowerShellCommand, String> {
22    let trimmed = command.trim();
23    if trimmed.is_empty() {
24        return Ok(ParsedPowerShellCommand { commands: Vec::new() });
25    }
26    let mut parser = tree_sitter::Parser::new();
27    parser
28        .set_language(&tree_sitter_pwsh::LANGUAGE.into())
29        .map_err(|err| format!("failed to load PowerShell grammar: {err}"))?;
30    let tree = parser
31        .parse(trimmed, None)
32        .ok_or_else(|| "PowerShell parser returned no tree".to_string())?;
33    if tree.root_node().has_error() {
34        return Err("PowerShell parse error".into());
35    }
36    Ok(ParsedPowerShellCommand { commands: split_commands(trimmed) })
37}
38
39fn split_commands(command: &str) -> Vec<ParsedCommandElement> {
40    command
41        .split([';', '\n'])
42        .flat_map(|part| part.split('|'))
43        .filter_map(parse_command_part)
44        .collect()
45}
46
47fn parse_command_part(part: &str) -> Option<ParsedCommandElement> {
48    let part = part.trim();
49    if part.is_empty() {
50        return None;
51    }
52    let dynamic = part.starts_with('&')
53        && !part
54            .trim_start_matches('&')
55            .trim_start()
56            .chars()
57            .next()
58            .map(|ch| ch.is_ascii_alphabetic() || ch == '_' || ch == '.')
59            .unwrap_or(false);
60    let part = part.trim_start_matches('&').trim_start();
61    let tokens = tokenize_words(part);
62    let name = tokens.first()?.clone();
63    let args = tokens.into_iter().skip(1).collect();
64    Some(ParsedCommandElement { name, args, dynamic })
65}
66
67fn tokenize_words(input: &str) -> Vec<String> {
68    let mut tokens = Vec::new();
69    let mut current = String::new();
70    let mut quote: Option<char> = None;
71    for ch in input.chars() {
72        if let Some(q) = quote {
73            if ch == q {
74                quote = None;
75            } else {
76                current.push(ch);
77            }
78            continue;
79        }
80        match ch {
81            '\'' | '"' => quote = Some(ch),
82            ch if ch.is_whitespace() => {
83                if !current.is_empty() {
84                    tokens.push(std::mem::take(&mut current));
85                }
86            }
87            _ => current.push(ch),
88        }
89    }
90    if !current.is_empty() {
91        tokens.push(current);
92    }
93    tokens
94}
95
96#[cfg(test)]
97mod tests {
98    use super::*;
99
100    #[test]
101    fn parses_simple_command_name_and_args() {
102        let parsed = parse("Get-Process -Name pwsh").expect("parse should succeed");
103        let commands = parsed.commands();
104        assert_eq!(commands.len(), 1);
105        assert_eq!(commands[0].name, "Get-Process");
106        assert_eq!(commands[0].args, vec!["-Name", "pwsh"]);
107    }
108
109    #[test]
110    fn marks_dynamic_invocation_as_dynamic() {
111        let parsed = parse("& ('i' + 'ex') 'payload'").expect("parse should succeed");
112        assert!(parsed.commands().iter().any(|cmd| cmd.dynamic));
113    }
114
115    #[test]
116    fn parse_failure_is_reported() {
117        let parsed = parse("Get-Process |");
118        assert!(parsed.is_err());
119    }
120}
telos_agent/tools/command_security/powershell/parser.rs

telos_agent/tools/command_security/powershell/
parser.rs
