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