ps_parser/parser/
token.rs

1use std::{collections::BTreeSet, fmt::Display};
2
3use super::script_result::PsValue;
4
5/// Represents a parsed PowerShell method call token.
6///
7/// Stores the original token string, the method name, and its arguments as
8/// `PsValue`s. Useful for analyzing and reconstructing method calls in scripts.
9#[derive(Debug, Clone, PartialEq)]
10pub struct MethodToken {
11    token: String,
12    self_: PsValue,
13    name: String,
14    arguments: Vec<PsValue>,
15}
16
17impl MethodToken {
18    pub fn new(token: String, self_: PsValue, name: String, arguments: Vec<PsValue>) -> Self {
19        Self {
20            token,
21            self_,
22            name,
23            arguments,
24        }
25    }
26
27    pub fn token(&self) -> &String {
28        &self.token
29    }
30
31    pub fn name(&self) -> &String {
32        &self.name
33    }
34
35    pub fn self_(&self) -> &PsValue {
36        &self.self_
37    }
38
39    pub fn args(&self) -> &Vec<PsValue> {
40        &self.arguments
41    }
42}
43
44#[derive(Debug, Clone, PartialEq)]
45pub struct CommandToken {
46    token: String,
47    name: String,
48    arguments: Vec<String>,
49}
50
51/// Represents a parsed PowerShell command token.
52///
53/// Stores the original token string, the command name, and its arguments as
54/// strings. Useful for identifying and reconstructing command invocations.
55impl CommandToken {
56    pub fn new(token: String, name: String, arguments: Vec<String>) -> Self {
57        Self {
58            token,
59            name,
60            arguments,
61        }
62    }
63
64    pub fn token(&self) -> &String {
65        &self.token
66    }
67
68    pub fn name(&self) -> &String {
69        &self.name
70    }
71
72    pub fn args(&self) -> &Vec<String> {
73        &self.arguments
74    }
75}
76
77/// Represents a parsed PowerShell expression token.
78///
79/// Stores the original token string and its evaluated value as `PsValue`.
80/// Useful for deobfuscation and analysis of expressions.
81#[derive(Debug, Clone, PartialEq)]
82pub struct ExpressionToken {
83    token: String,
84    value: PsValue,
85}
86
87impl ExpressionToken {
88    pub fn new(token: String, value: PsValue) -> Self {
89        Self { token, value }
90    }
91}
92
93/// Represents a double-quoted PowerShell string with variable expansion.
94///
95/// Stores the original token string and its expanded value.
96/// Useful for tracking and reconstructing expandable strings in scripts.
97#[derive(Debug, Clone, PartialEq)]
98pub struct StringExpandableToken {
99    token: String,
100    value: String,
101}
102
103impl StringExpandableToken {
104    pub fn new(token: String, value: String) -> Self {
105        Self { token, value }
106    }
107}
108
109#[derive(Debug, Clone, PartialEq)]
110pub enum Token {
111    StringExpandable(StringExpandableToken),
112    String(String),
113    Expression(ExpressionToken),
114    Method(MethodToken),
115    Command(CommandToken),
116}
117impl Token {
118    pub fn method(token: String, self_: PsValue, name: String, arguments: Vec<PsValue>) -> Self {
119        Token::Method(MethodToken {
120            token,
121            self_,
122            name,
123            arguments,
124        })
125    }
126
127    pub fn command(token: String, name: String, arguments: Vec<String>) -> Self {
128        Token::Command(CommandToken {
129            token,
130            name,
131            arguments,
132        })
133    }
134
135    pub fn expression(token: String, value: PsValue) -> Self {
136        Token::Expression(ExpressionToken { token, value })
137    }
138
139    pub fn string_expandable(token: String, value: String) -> Self {
140        Token::StringExpandable(StringExpandableToken { token, value })
141    }
142}
143impl Display for Token {
144    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
145        write!(f, "{:?}", self)
146    }
147}
148
149#[derive(Default, Debug, Clone, PartialEq)]
150pub struct Tokens(Vec<Token>);
151impl Tokens {
152    pub fn new() -> Self {
153        Self(Vec::new())
154    }
155
156    pub fn push(&mut self, token: Token) {
157        self.0.push(token)
158    }
159
160    pub fn all(&self) -> Vec<Token> {
161        self.0.clone()
162    }
163
164    pub fn literal_strings(&self) -> Vec<String> {
165        self.0
166            .iter()
167            .filter_map(|token| match token {
168                Token::String(literal) => Some(literal.clone()),
169                _ => None,
170            })
171            .collect()
172    }
173
174    pub fn string_set(&self) -> BTreeSet<String> {
175        let mut string_set = BTreeSet::new();
176        for token in self.0.iter() {
177            match token {
178                Token::String(deobfuscated)
179                | Token::StringExpandable(StringExpandableToken {
180                    value: deobfuscated,
181                    ..
182                }) => {
183                    let _ = string_set.insert(deobfuscated.to_string());
184                }
185                _ => {}
186            }
187        }
188        string_set
189    }
190
191    pub fn lowercased_string_set(&self) -> BTreeSet<String> {
192        let mut string_set = BTreeSet::new();
193        for token in self.0.iter() {
194            match token {
195                Token::String(deobfuscated)
196                | Token::StringExpandable(StringExpandableToken {
197                    value: deobfuscated,
198                    ..
199                }) => {
200                    let _ = string_set.insert(deobfuscated.to_ascii_lowercase());
201                }
202                _ => {}
203            }
204        }
205        string_set
206    }
207
208    pub fn expandable_strings(&self) -> Vec<StringExpandableToken> {
209        self.0
210            .iter()
211            .filter_map(|token| match token {
212                Token::StringExpandable(expandable) => Some(expandable.clone()),
213                _ => None,
214            })
215            .collect()
216    }
217
218    pub fn expressions(&self) -> Vec<ExpressionToken> {
219        self.0
220            .iter()
221            .filter_map(|token| match token {
222                Token::Expression(expr) => Some(expr.clone()),
223                _ => None,
224            })
225            .collect()
226    }
227
228    pub fn methods(&self) -> Vec<MethodToken> {
229        self.0
230            .iter()
231            .filter_map(|token| match token {
232                Token::Method(method) => Some(method.clone()),
233                _ => None,
234            })
235            .collect()
236    }
237
238    pub fn commands(&self) -> Vec<CommandToken> {
239        self.0
240            .iter()
241            .filter_map(|token| match token {
242                Token::Command(command) => Some(command.clone()),
243                _ => None,
244            })
245            .collect()
246    }
247}
ps_parser/parser/token.rs

ps_parser/parser/
token.rs
