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    pub fn token(&self) -> &String {
109        &self.token
110    }
111
112    pub fn value(&self) -> &String {
113        &self.value
114    }
115}
116
117#[derive(Debug, Clone, PartialEq)]
118pub struct TypeToken {
119    pub name: String,
120}
121impl TypeToken {
122    pub fn new(name: String) -> Self {
123        Self { name }
124    }
125
126    pub fn name(&self) -> &String {
127        &self.name
128    }
129}
130
131#[derive(Debug, Clone, PartialEq)]
132pub enum Token {
133    StringExpandable(StringExpandableToken),
134    String(String),
135    Expression(ExpressionToken),
136    Method(MethodToken),
137    Command(CommandToken),
138    Type(TypeToken),
139}
140impl Token {
141    pub fn method(token: String, self_: PsValue, name: String, arguments: Vec<PsValue>) -> Self {
142        Token::Method(MethodToken {
143            token,
144            self_,
145            name,
146            arguments,
147        })
148    }
149
150    pub fn command(token: String, name: String, arguments: Vec<String>) -> Self {
151        Token::Command(CommandToken {
152            token,
153            name,
154            arguments,
155        })
156    }
157
158    pub fn expression(token: String, value: PsValue) -> Self {
159        Token::Expression(ExpressionToken { token, value })
160    }
161
162    pub fn string_expandable(token: String, value: String) -> Self {
163        Token::StringExpandable(StringExpandableToken { token, value })
164    }
165
166    pub fn type_literal(name: String) -> Self {
167        Token::Type(TypeToken { name })
168    }
169}
170impl Display for Token {
171    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
172        write!(f, "{:?}", self)
173    }
174}
175
176#[derive(Default, Debug, Clone, PartialEq)]
177pub struct Tokens(Vec<Token>);
178impl Tokens {
179    pub fn new() -> Self {
180        Self(Vec::new())
181    }
182
183    pub fn push(&mut self, token: Token) {
184        self.0.push(token)
185    }
186
187    pub fn all(&self) -> Vec<Token> {
188        self.0.clone()
189    }
190
191    pub fn literal_strings(&self) -> Vec<String> {
192        self.0
193            .iter()
194            .filter_map(|token| match token {
195                Token::String(literal) => Some(literal.clone()),
196                _ => None,
197            })
198            .collect()
199    }
200
201    pub fn string_set(&self) -> BTreeSet<String> {
202        let mut string_set = BTreeSet::new();
203        for token in self.0.iter() {
204            match token {
205                Token::String(deobfuscated)
206                | Token::StringExpandable(StringExpandableToken {
207                    value: deobfuscated,
208                    ..
209                }) => {
210                    let _ = string_set.insert(deobfuscated.to_string());
211                }
212                _ => {}
213            }
214        }
215        string_set
216    }
217
218    pub fn lowercased_string_set(&self) -> BTreeSet<String> {
219        let mut string_set = BTreeSet::new();
220        for token in self.0.iter() {
221            match token {
222                Token::String(deobfuscated)
223                | Token::StringExpandable(StringExpandableToken {
224                    value: deobfuscated,
225                    ..
226                }) => {
227                    let _ = string_set.insert(deobfuscated.to_ascii_lowercase());
228                }
229                _ => {}
230            }
231        }
232        string_set
233    }
234
235    pub fn expandable_strings(&self) -> Vec<StringExpandableToken> {
236        self.0
237            .iter()
238            .filter_map(|token| match token {
239                Token::StringExpandable(expandable) => Some(expandable.clone()),
240                _ => None,
241            })
242            .collect()
243    }
244
245    pub fn expressions(&self) -> Vec<ExpressionToken> {
246        self.0
247            .iter()
248            .filter_map(|token| match token {
249                Token::Expression(expr) => Some(expr.clone()),
250                _ => None,
251            })
252            .collect()
253    }
254
255    pub fn methods(&self) -> Vec<MethodToken> {
256        self.0
257            .iter()
258            .filter_map(|token| match token {
259                Token::Method(method) => Some(method.clone()),
260                _ => None,
261            })
262            .collect()
263    }
264
265    pub fn commands(&self) -> Vec<CommandToken> {
266        self.0
267            .iter()
268            .filter_map(|token| match token {
269                Token::Command(command) => Some(command.clone()),
270                _ => None,
271            })
272            .collect()
273    }
274
275    pub fn methods_and_commands(&self) -> BTreeSet<&str> {
276        self.0
277            .iter()
278            .filter_map(|token| match token {
279                Token::Command(command) => Some(command.name().as_str()),
280                Token::Method(method) => Some(method.name().as_str()),
281                _ => None,
282            })
283            .collect()
284    }
285
286    pub fn ttypes(&self) -> BTreeSet<&str> {
287        self.0
288            .iter()
289            .filter_map(|token| match token {
290                Token::Type(type_token) => Some(type_token.name().as_str()),
291                _ => None,
292            })
293            .collect()
294    }
295}