1use std::{collections::BTreeSet, fmt::Display};
2
3use super::script_result::PsValue;
4
5#[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
51impl 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#[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#[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}