1mod parser;
30pub(crate) use parser::NEWLINE;
31pub use parser::PowerShellSession;
54pub use parser::PsValue;
70pub use parser::ScriptResult;
90pub use parser::Token;
122pub use parser::Variables;
147
148#[cfg(test)]
149mod tests {
150 use std::collections::HashMap;
151
152 use super::*;
153 use crate::Token;
154
155 #[test]
156 fn deobfuscation() {
157 let mut p = PowerShellSession::new();
159 let input = r#" $script:var = [char]([int]("9e4e" -replace "e")+3); [int]'a';$var"#;
160 let script_res = p.parse_input(input).unwrap();
161 assert_eq!(script_res.result(), 'a'.into());
162 assert_eq!(
163 script_res.deobfuscated(),
164 vec!["$script:var = 'a'", "[int]'a'", "'a'"].join(NEWLINE)
165 );
166 assert_eq!(script_res.errors().len(), 1);
167 assert_eq!(
168 script_res.errors()[0].to_string(),
169 "ValError: Cannot convert value \"String\" to type \"Int\""
170 );
171
172 let mut p = PowerShellSession::new();
174 let input = r#" $global:var = [char]([int]("9e4e" -replace "e")+3) "#;
175 let script_res = p.parse_input(input).unwrap();
176
177 assert_eq!(script_res.errors().len(), 0);
178
179 let script_res = p.parse_input(" [int]'a';$var ").unwrap();
180 assert_eq!(
181 script_res.deobfuscated(),
182 vec!["[int]'a'", "'a'"].join(NEWLINE)
183 );
184 assert_eq!(script_res.output(), vec!["a"].join(NEWLINE));
185 assert_eq!(script_res.errors().len(), 1);
186 assert_eq!(
187 script_res.errors()[0].to_string(),
188 "ValError: Cannot convert value \"String\" to type \"Int\""
189 );
190 }
191
192 #[test]
193 fn deobfuscation_non_existing_value() {
194 let mut p = PowerShellSession::new();
196 let input = r#" $local:var = $env:programfiles;[int]'a';$var"#;
197 let script_res = p.parse_input(input).unwrap();
198 assert_eq!(script_res.result(), PsValue::Null);
199 assert_eq!(
200 script_res.deobfuscated(),
201 vec!["$local:var = $env:programfiles", "[int]'a'", "$var"].join(NEWLINE)
202 );
203 assert_eq!(script_res.errors().len(), 3);
204 assert_eq!(
205 script_res.errors()[0].to_string(),
206 "VariableError: Variable \"programfiles\" is not defined"
207 );
208 assert_eq!(
209 script_res.errors()[1].to_string(),
210 "ValError: Cannot convert value \"String\" to type \"Int\""
211 );
212 assert_eq!(
213 script_res.errors()[2].to_string(),
214 "VariableError: Variable \"var\" is not defined"
215 );
216
217 let mut p = PowerShellSession::new().with_variables(Variables::force_eval());
219 let input = r#" $local:var = $env:programfiles;[int]'a';$script:var"#;
220 let script_res = p.parse_input(input).unwrap();
221 assert_eq!(script_res.result(), PsValue::Null);
222 assert_eq!(
223 script_res.deobfuscated(),
224 vec!["$local:var = $null", "[int]'a'"].join(NEWLINE)
225 );
226 assert_eq!(script_res.errors().len(), 1);
227 }
228
229 #[test]
230 fn deobfuscation_env_value() {
231 let mut p = PowerShellSession::new().with_variables(Variables::env());
233 let input = r#" $local:var = $env:programfiles;$var"#;
234 let script_res = p.parse_input(input).unwrap();
235 assert_eq!(
236 script_res.result(),
237 PsValue::String(std::env::var("PROGRAMFILES").unwrap())
238 );
239 let program_files = std::env::var("PROGRAMFILES").unwrap();
240 assert_eq!(
241 script_res.deobfuscated(),
242 vec![
243 format!("$local:var = \"{}\"", program_files),
244 format!("\"{}\"", program_files)
245 ]
246 .join(NEWLINE)
247 );
248 assert_eq!(script_res.errors().len(), 0);
249 }
250
251 #[test]
252 fn hash_table() {
253 let mut p = PowerShellSession::new().with_variables(Variables::env().values_persist());
255 let input = r#"
256$nestedData = @{
257 Users = @(
258 @{ Name = "Alice"; Age = 30; Skills = @("PowerShell", "Python") }
259 @{ Name = "Bob"; Age = 25; Skills = @("Java", "C#") }
260 )
261 Settings = @{
262 Theme = "Dark"
263 Language = "en-US"
264 }
265}
266"$nestedData"
267 "#;
268 let script_res = p.parse_input(input).unwrap();
269 assert_eq!(
270 script_res.result(),
271 PsValue::String("System.Collections.Hashtable".to_string())
272 );
273
274 assert_eq!(
275 p.parse_input("$nesteddata.settings").unwrap().result(),
276 PsValue::HashTable(HashMap::from([
277 ("language".to_string(), PsValue::String("en-US".to_string())),
278 ("theme".to_string(), PsValue::String("Dark".to_string())),
279 ]))
280 );
281
282 assert_eq!(
283 p.safe_eval("$nesteddata.settings.theme").unwrap(),
284 "Dark".to_string()
285 );
286
287 assert_eq!(
288 p.parse_input("$nesteddata.users[0]").unwrap().result(),
289 PsValue::HashTable(HashMap::from([
290 (
291 "skills".to_string(),
292 PsValue::Array(vec![
293 PsValue::String("PowerShell".to_string()),
294 PsValue::String("Python".to_string().into())
295 ])
296 ),
297 ("name".to_string(), PsValue::String("Alice".to_string())),
298 ("age".to_string(), PsValue::Int(30)),
299 ]))
300 );
301
302 assert_eq!(
303 p.safe_eval("$nesteddata.users[0]['name']").unwrap(),
304 "Alice".to_string()
305 );
306
307 assert_eq!(
308 p.safe_eval("$nesteddata.users[0].NAME").unwrap(),
309 "Alice".to_string()
310 );
311 }
312
313 #[test]
314 fn test_simple_arithmetic() {
315 let input = r#"
316Write-Host "=== Test 3: Arithmetic Operations ===" -ForegroundColor Green
317$a = 10
318$b = 5
319Write-Output "Addition: $(($a + $b))"
320Write-Output "Subtraction: $(($a - $b))"
321Write-Output "Multiplication: $(($a * $b))"
322Write-Output "Division: $(($a / $b))"
323Write-Output "Modulo: $(($a % $b))"
324"#;
325
326 let script_result = PowerShellSession::new().parse_input(input).unwrap();
327
328 assert_eq!(script_result.result(), PsValue::String("Modulo: 0".into()));
329 assert_eq!(
330 script_result.output(),
331 vec![
332 r#"=== Test 3: Arithmetic Operations ==="#,
333 r#"Addition: 15"#,
334 r#"Subtraction: 5"#,
335 r#"Multiplication: 50"#,
336 r#"Division: 2"#,
337 r#"Modulo: 0"#
338 ]
339 .join(NEWLINE)
340 );
341 assert_eq!(script_result.errors().len(), 0);
342 assert_eq!(script_result.tokens().strings(), vec![]);
343 assert_eq!(script_result.tokens().expandable_strings().len(), 6);
344 assert_eq!(
345 script_result.tokens().expandable_strings()[1],
346 Token::StringExpandable(
347 "\"Addition: $(($a + $b))\"".to_string(),
348 "Addition: 15".to_string()
349 )
350 );
351 assert_eq!(script_result.tokens().expression().len(), 12);
352 assert_eq!(
353 script_result.tokens().expression()[2],
354 Token::Expression("$a + $b".to_string(), PsValue::Int(15))
355 );
356 }
357
358 #[test]
359 fn test_scripts() {
360 use std::fs;
361 let Ok(entries) = fs::read_dir("test_scripts") else {
362 panic!("Failed to read 'test_scripts' directory");
363 };
364 for entry in entries {
365 let dir_entry = entry.unwrap();
366 if std::fs::FileType::is_dir(&dir_entry.file_type().unwrap()) {
367 let input_script = dir_entry.path().join("input.ps1");
369 let expected_deobfuscated_script = dir_entry.path().join("deobfuscated.txt");
370 let expected_output_script = dir_entry.path().join("output.txt");
371
372 let Ok(input) = fs::read_to_string(&input_script) else {
373 panic!("Failed to read test file: {}", input_script.display());
374 };
375
376 let Ok(expected_deobfuscated) = fs::read_to_string(&expected_deobfuscated_script)
377 else {
378 panic!(
379 "Failed to read test file: {}",
380 expected_deobfuscated_script.display()
381 );
382 };
383
384 let Ok(expected_output) = fs::read_to_string(&expected_output_script) else {
385 panic!(
386 "Failed to read test file: {}",
387 expected_output_script.display()
388 );
389 };
390
391 let script_result = PowerShellSession::new()
392 .with_variables(Variables::env())
393 .parse_input(&input)
394 .unwrap();
395
396 let expected_deobfuscated_vec = expected_deobfuscated
397 .lines()
398 .map(|s| s.trim_end())
399 .collect::<Vec<&str>>();
400
401 let current_deobfuscated = script_result.deobfuscated();
402 let current_output = script_result.output();
403
404 let expected_output_vec = expected_output
405 .lines()
406 .map(|s| s.trim_end())
407 .collect::<Vec<&str>>();
408
409 let current_deobfuscated_vec = current_deobfuscated
411 .lines()
412 .map(|s| s.trim_end())
413 .collect::<Vec<&str>>();
414
415 let current_output_vec = current_output
416 .lines()
417 .map(|s| s.trim_end())
418 .collect::<Vec<&str>>();
419
420 for i in 0..expected_deobfuscated_vec.len() {
421 assert_eq!(
422 expected_deobfuscated_vec[i],
423 current_deobfuscated_vec[i],
424 "File: {}, Deobfuscated line: {}",
425 file_name(&dir_entry),
426 i + 1
427 );
428 }
429
430 for i in 0..expected_output_vec.len() {
431 assert_eq!(
432 expected_output_vec[i],
433 current_output_vec[i],
434 "File: {}, Output line: {}",
435 file_name(&dir_entry),
436 i + 1
437 );
438 }
439 }
440 }
441 }
442
443 fn file_name(dir_entry: &std::fs::DirEntry) -> String {
444 dir_entry
445 .path()
446 .components()
447 .last()
448 .unwrap()
449 .as_os_str()
450 .to_string_lossy()
451 .to_string()
452 }
453
454 #[allow(dead_code)]
455 fn save_files(dir_entry: &std::fs::DirEntry, deobfuscated: &str, output: &str) {
456 let name = file_name(dir_entry);
457 std::fs::write(format!("{}_deobfuscated.txt", name), deobfuscated).unwrap();
458 std::fs::write(format!("{}_output.txt", name), output).unwrap();
459 }
460
461 #[test]
462 fn test_range() {
463 let mut p = PowerShellSession::new().with_variables(Variables::env());
465 let input = r#" $numbers = 1..10; $numbers"#;
466 let script_res = p.parse_input(input).unwrap();
467 assert_eq!(
468 script_res.deobfuscated(),
469 vec![
470 "$numbers = @(1,2,3,4,5,6,7,8,9,10)",
471 "@(1,2,3,4,5,6,7,8,9,10)"
472 ]
473 .join(NEWLINE)
474 );
475 assert_eq!(script_res.errors().len(), 0);
476 }
477
478 #[test]
479 fn even_numbers() {
480 let mut p = PowerShellSession::new().with_variables(Variables::env());
482 let input = r#" $numbers = 1..10; $evenNumbers = $numbers | Where-Object { $_ % 2 -eq 0 }; $evenNumbers"#;
483 let script_res = p.parse_input(input).unwrap();
484 assert_eq!(
485 script_res.result(),
486 PsValue::Array(vec![
487 PsValue::Int(2),
488 PsValue::Int(4),
489 PsValue::Int(6),
490 PsValue::Int(8),
491 PsValue::Int(10)
492 ])
493 );
494 assert_eq!(
495 script_res.deobfuscated(),
496 vec![
497 "$numbers = @(1,2,3,4,5,6,7,8,9,10)",
498 "$evennumbers = @(2,4,6,8,10)",
499 "@(2,4,6,8,10)"
500 ]
501 .join(NEWLINE)
502 );
503 assert_eq!(script_res.errors().len(), 0);
504 }
505
506 #[test]
507 fn divisible_by_2_and_3() {
508 let mut p = PowerShellSession::new().with_variables(Variables::env());
510 let input = r#" $numbers = 1..10; $numbers | Where { $_ % 2 -eq 0 } | ? { $_ % 3 -eq 0 }"#;
511 let script_res = p.parse_input(input).unwrap();
512 assert_eq!(script_res.result(), PsValue::Array(vec![PsValue::Int(6),]));
513 assert_eq!(
514 script_res.deobfuscated(),
515 vec!["$numbers = @(1,2,3,4,5,6,7,8,9,10)", "@(6)"].join(NEWLINE)
516 );
517 assert_eq!(script_res.errors().len(), 0);
518 }
519
520 fn _test_function() {
522 let mut p = PowerShellSession::new().with_variables(Variables::env());
524 let input = r#"
525function Get-Square($number) {
526 return $number * $number
527}
528"Square of 5: $(Get-Square 5)" "#;
529 let script_res = p.parse_input(input).unwrap();
530 assert_eq!(
531 script_res.deobfuscated(),
532 vec![
533 "function Get-Square($number) {",
534 " return $number * $number",
535 "}",
536 " \"Square of 5: $(Get-Square 5)\""
537 ]
538 .join(NEWLINE)
539 );
540 assert_eq!(script_res.errors().len(), 2);
541 }
542
543 #[test]
544 fn test_if() {
545 let mut p = PowerShellSession::new().with_variables(Variables::env());
547 let input = r#"
548 # Test 10: Conditional Statements
549if ($true) {
550 $if_result = "condition true"
551}
552
553if ($false) {
554 $else_result = "false branch"
555} else {
556 $else_result = "true branch"
557}
558
559$score = 85
560if ($score -ge 90) {
561 $grade = "A"
562} elseif ($score -ge 80) {
563 $grade = "B"
564} else {
565 $grade = "C"
566}
567
568 "#;
569 let script_res = p.parse_input(input).unwrap();
570 assert_eq!(
571 script_res.deobfuscated(),
572 vec![
573 "$if_result = \"condition true\"",
574 "$else_result = \"true branch\"",
575 "$score = 85",
576 "$grade = \"B\""
577 ]
578 .join(NEWLINE)
579 );
580 assert_eq!(script_res.errors().len(), 0);
581 }
582}