1mod function;
2mod scopes;
3mod variable;
4
5use std::collections::HashMap;
6
7use function::FunctionMap;
8use phf::phf_map;
9pub(super) use scopes::SessionScope;
10use thiserror_no_std::Error;
11pub(super) use variable::{Scope, VarName};
12
13use crate::parser::{Val, value::ScriptBlock};
14#[derive(Error, Debug, PartialEq, Clone)]
15pub enum VariableError {
16 #[error("Variable \"{0}\" is not defined")]
17 NotDefined(String),
18 #[error("Cannot overwrite variable \"{0}\" because it is read-only or constant.")]
19 ReadOnly(String),
20}
21
22pub type VariableResult<T> = core::result::Result<T, VariableError>;
23pub type VariableMap = HashMap<String, Val>;
24
25#[derive(Clone, Default)]
26pub struct Variables {
27 env: VariableMap,
28 global_scope: VariableMap,
29 script_scope: VariableMap,
30 scope_sessions_stack: Vec<VariableMap>,
31 state: State,
32 force_var_eval: bool,
33 values_persist: bool,
34 global_functions: FunctionMap,
35 script_functions: FunctionMap,
36 }
42
43#[derive(Default, Clone)]
44enum State {
45 #[default]
46 Script,
47 Stack(u32),
48}
49
50impl Variables {
51 const PREDEFINED_VARIABLES: phf::Map<&'static str, Val> = phf_map! {
52 "true" => Val::Bool(true),
53 "false" => Val::Bool(false),
54 "null" => Val::Null,
55 };
56
57 pub(crate) fn set_ps_item(&mut self, ps_item: Val) {
58 let _ = self.set(
59 &VarName::new_with_scope(Scope::Special, "$PSItem".into()),
60 ps_item.clone(),
61 );
62 let _ = self.set(
63 &VarName::new_with_scope(Scope::Special, "$_".into()),
64 ps_item,
65 );
66 }
67
68 pub(crate) fn reset_ps_item(&mut self) {
69 let _ = self.set(
70 &VarName::new_with_scope(Scope::Special, "$PSItem".into()),
71 Val::Null,
72 );
73 let _ = self.set(
74 &VarName::new_with_scope(Scope::Special, "$_".into()),
75 Val::Null,
76 );
77 }
78
79 pub fn set_status(&mut self, b: bool) {
80 let _ = self.set(
81 &VarName::new_with_scope(Scope::Special, "$?".into()),
82 Val::Bool(b),
83 );
84 }
85
86 pub fn status(&mut self) -> bool {
87 let Some(Val::Bool(b)) = self.get(&VarName::new_with_scope(Scope::Special, "$?".into()))
88 else {
89 return false;
90 };
91 b
92 }
93
94 pub fn load_from_file(
95 &mut self,
96 path: &std::path::Path,
97 ) -> Result<(), Box<dyn std::error::Error>> {
98 let mut config_parser = configparser::ini::Ini::new();
99 let map = config_parser.load(path)?;
100 self.load(map)
101 }
102
103 pub fn load_from_string(&mut self, ini_string: &str) -> Result<(), Box<dyn std::error::Error>> {
104 let mut config_parser = configparser::ini::Ini::new();
105 let map = config_parser.read(ini_string.into())?;
106 self.load(map)
107 }
108
109 pub fn init(&mut self) {
110 if !self.values_persist {
111 self.script_scope.clear();
112 }
113 self.scope_sessions_stack.clear();
114 self.state = State::Script;
115 }
116
117 fn load(
118 &mut self,
119 conf_map: HashMap<String, HashMap<String, Option<String>>>,
120 ) -> Result<(), Box<dyn std::error::Error>> {
121 for (section_name, properties) in conf_map {
122 for (key, value) in properties {
123 let Some(value) = value else {
124 continue;
125 };
126
127 let var_name = match section_name.as_str() {
128 "global" => VarName::new_with_scope(Scope::Global, key.to_lowercase()),
129 "script" => VarName::new_with_scope(Scope::Script, key.to_lowercase()),
130 "env" => VarName::new_with_scope(Scope::Env, key.to_lowercase()),
131 _ => {
132 continue;
133 }
134 };
135
136 let parsed_value = if let Ok(bool_val) = value.parse::<bool>() {
138 Val::Bool(bool_val)
139 } else if let Ok(int_val) = value.parse::<i64>() {
140 Val::Int(int_val)
141 } else if let Ok(float_val) = value.parse::<f64>() {
142 Val::Float(float_val)
143 } else if value.is_empty() {
144 Val::Null
145 } else {
146 Val::String(value.clone().into())
147 };
148
149 if let Err(err) = self.set(&var_name, parsed_value.clone()) {
151 log::error!("Failed to set variable {:?}: {}", var_name, err);
152 }
153 }
154 }
155 Ok(())
156 }
157
158 pub(crate) fn script_scope(&self) -> VariableMap {
159 self.script_scope.clone()
160 }
161
162 pub(crate) fn get_env(&self) -> VariableMap {
163 self.env.clone()
164 }
165
166 pub(crate) fn get_global(&self) -> VariableMap {
167 self.global_scope.clone()
168 }
169
170 pub(crate) fn add_script_function(&mut self, name: String, func: ScriptBlock) {
171 self.script_functions.insert(name, func);
172 }
173
174 pub(crate) fn add_global_function(&mut self, name: String, func: ScriptBlock) {
175 self.global_functions.insert(name, func);
176 }
177
178 pub(crate) fn clear_script_functions(&mut self) {
179 self.script_functions.clear();
180 }
181
182 pub fn new() -> Variables {
205 Default::default()
206 }
207
208 pub fn force_eval() -> Self {
243 Self {
244 force_var_eval: true,
245 ..Default::default()
246 }
247 }
248
249 #[allow(dead_code)]
251 pub(crate) fn values_persist(mut self) -> Self {
252 self.values_persist = true;
253 self
254 }
255
256 pub fn env() -> Variables {
279 let mut vars = Variables::new();
280
281 for (key, value) in std::env::vars() {
283 vars.env
286 .insert(key.to_lowercase(), Val::String(value.into()));
287 }
288 vars
289 }
290
291 pub fn from_ini_string(ini_string: &str) -> Result<Self, Box<dyn std::error::Error>> {
336 let mut variables = Self::new();
337 variables.load_from_string(ini_string)?;
338 Ok(variables)
339 }
340
341 pub fn from_ini_file(path: &std::path::Path) -> Result<Self, Box<dyn std::error::Error>> {
343 let mut variables = Self::new();
344 variables.load_from_file(path)?;
345 Ok(variables)
346 }
347
348 fn const_map_from_scope(&self, scope: &Scope) -> &VariableMap {
349 match scope {
350 Scope::Global => &self.global_scope,
351 Scope::Script => &self.script_scope,
352 Scope::Env => &self.env,
353 Scope::Local => match self.state {
354 State::Script => &self.script_scope,
355 State::Stack(depth) => {
356 if depth < self.scope_sessions_stack.len() as u32 {
357 &self.scope_sessions_stack[depth as usize]
358 } else {
359 &self.script_scope
360 }
361 }
362 },
363 Scope::Special => {
364 &self.global_scope }
366 }
367 }
368
369 fn local_scope(&mut self) -> &mut VariableMap {
370 match self.state {
371 State::Script => &mut self.script_scope,
372 State::Stack(depth) => {
373 if depth < self.scope_sessions_stack.len() as u32 {
374 &mut self.scope_sessions_stack[depth as usize]
375 } else {
376 &mut self.script_scope
377 }
378 }
379 }
380 }
381 fn map_from_scope(&mut self, scope: &Scope) -> &mut VariableMap {
382 match scope {
383 Scope::Global => &mut self.global_scope,
384 Scope::Script => &mut self.script_scope,
385 Scope::Env => &mut self.env,
386 Scope::Local => self.local_scope(),
387 Scope::Special => {
388 &mut self.global_scope }
390 }
391 }
392
393 pub(crate) fn set(&mut self, var_name: &VarName, val: Val) -> VariableResult<()> {
405 let var = self.find_mut_variable_in_scopes(var_name)?;
406
407 if let Some(variable) = var {
408 *variable = val;
409 } else {
410 let map = self.map_from_scope(&var_name.scope.clone().unwrap_or(Scope::Local));
411 map.insert(var_name.name.to_ascii_lowercase(), val);
412 }
413
414 Ok(())
415 }
416
417 pub(crate) fn set_local(&mut self, name: &str, val: Val) -> VariableResult<()> {
418 let var_name = VarName::new_with_scope(Scope::Local, name.to_ascii_lowercase());
419 self.set(&var_name, val)
420 }
421
422 fn find_mut_variable_in_scopes(
423 &mut self,
424 var_name: &VarName,
425 ) -> VariableResult<Option<&mut Val>> {
426 let name = var_name.name.to_ascii_lowercase();
427 let name_str = name.as_str();
428
429 if let Some(scope) = &var_name.scope {
430 let map = self.map_from_scope(scope);
431 Ok(map.get_mut(name_str))
432 } else {
433 if Self::PREDEFINED_VARIABLES.contains_key(name_str) {
434 return Err(VariableError::ReadOnly(name.clone()));
435 }
436
437 for local_scope in self.scope_sessions_stack.iter_mut().rev() {
439 if local_scope.contains_key(name_str) {
440 return Ok(local_scope.get_mut(name_str));
441 }
442 }
443
444 if self.script_scope.contains_key(name_str) {
445 return Ok(self.script_scope.get_mut(name_str));
446 }
447
448 if self.global_scope.contains_key(name_str) {
449 return Ok(self.global_scope.get_mut(name_str));
450 }
451
452 Ok(None)
453 }
454 }
455
456 pub(crate) fn get(&self, var_name: &VarName) -> Option<Val> {
467 let var = self.find_variable_in_scopes(var_name);
468
469 if self.force_var_eval && var.is_none() {
470 Some(Val::Null)
471 } else {
472 var.cloned()
473 }
474 }
475
476 fn find_variable_in_scopes(&self, var_name: &VarName) -> Option<&Val> {
477 let name = var_name.name.to_ascii_lowercase();
478 let name_str = name.as_str();
479
480 if let Some(scope) = &var_name.scope {
481 let map = self.const_map_from_scope(scope);
482 map.get(name_str)
483 } else {
484 if Self::PREDEFINED_VARIABLES.contains_key(name_str) {
485 return Self::PREDEFINED_VARIABLES.get(name_str);
486 }
487
488 for local_scope in self.scope_sessions_stack.iter().rev() {
490 if local_scope.contains_key(name_str) {
491 return local_scope.get(name_str);
492 }
493 }
494
495 if self.script_scope.contains_key(name_str) {
496 return self.script_scope.get(name_str);
497 }
498
499 if self.global_scope.contains_key(name_str) {
500 return self.global_scope.get(name_str);
501 }
502
503 None
504 }
505 }
506
507 pub(crate) fn push_scope_session(&mut self) {
508 let current_map = self.local_scope();
509 let new_map = current_map.clone();
510
511 self.scope_sessions_stack.push(new_map);
512 self.state = State::Stack(self.scope_sessions_stack.len() as u32 - 1);
513 }
514
515 pub(crate) fn pop_scope_session(&mut self) {
516 match self.scope_sessions_stack.len() {
517 0 => {} 1 => {
519 self.scope_sessions_stack.pop();
520 self.state = State::Script;
521 }
522 _ => {
523 self.scope_sessions_stack.pop();
524 self.state = State::Stack(self.scope_sessions_stack.len() as u32 - 1);
525 }
526 }
527 }
528}
529
530#[cfg(test)]
531mod tests {
532 use super::Variables;
533 use crate::{PowerShellSession, PsValue};
534
535 #[test]
536 fn test_builtin_variables() {
537 let mut p = PowerShellSession::new();
538 assert_eq!(p.safe_eval(r#" $true "#).unwrap().as_str(), "True");
539 assert_eq!(p.safe_eval(r#" $false "#).unwrap().as_str(), "False");
540 assert_eq!(p.safe_eval(r#" $null "#).unwrap().as_str(), "");
541 }
542
543 #[test]
544 fn test_builtint_objects() {
545 let mut p = PowerShellSession::new();
546 assert_eq!(
547 p.parse_input(r#" [system.convert]0 "#).unwrap().result(),
548 PsValue::Null
549 );
550 }
551
552 #[test]
553 fn test_env_variables() {
554 let v = Variables::env();
555 let mut p = PowerShellSession::new().with_variables(v);
556 assert_eq!(
557 p.safe_eval(r#" $env:path "#).unwrap().as_str(),
558 std::env::var("PATH").unwrap()
559 );
560 assert_eq!(
561 p.safe_eval(r#" $env:username "#).unwrap().as_str(),
562 std::env::var("USERNAME").unwrap()
563 );
564 assert_eq!(
565 p.safe_eval(r#" $env:tEMp "#).unwrap().as_str(),
566 std::env::var("TEMP").unwrap()
567 );
568 assert_eq!(
569 p.safe_eval(r#" $env:tMp "#).unwrap().as_str(),
570 std::env::var("TMP").unwrap()
571 );
572 assert_eq!(
573 p.safe_eval(r#" $env:cOmputername "#).unwrap().as_str(),
574 std::env::var("COMPUTERNAME").unwrap()
575 );
576 assert_eq!(
577 p.safe_eval(r#" $env:programfiles "#).unwrap().as_str(),
578 std::env::var("PROGRAMFILES").unwrap()
579 );
580 assert_eq!(
581 p.safe_eval(r#" $env:temp "#).unwrap().as_str(),
582 std::env::var("TEMP").unwrap()
583 );
584 assert_eq!(
585 p.safe_eval(r#" ${Env:ProgramFiles(x86)} "#)
586 .unwrap()
587 .as_str(),
588 std::env::var("ProgramFiles(x86)").unwrap()
589 );
590 let env_variables = p.env_variables();
591 assert_eq!(
592 env_variables.get("path").unwrap().to_string(),
593 std::env::var("PATH").unwrap()
594 );
595 assert_eq!(
596 env_variables.get("tmp").unwrap().to_string(),
597 std::env::var("TMP").unwrap()
598 );
599 assert_eq!(
600 env_variables.get("temp").unwrap().to_string(),
601 std::env::var("TMP").unwrap()
602 );
603 assert_eq!(
604 env_variables.get("appdata").unwrap().to_string(),
605 std::env::var("APPDATA").unwrap()
606 );
607 assert_eq!(
608 env_variables.get("username").unwrap().to_string(),
609 std::env::var("USERNAME").unwrap()
610 );
611 assert_eq!(
612 env_variables.get("programfiles").unwrap().to_string(),
613 std::env::var("PROGRAMFILES").unwrap()
614 );
615 assert_eq!(
616 env_variables.get("programfiles(x86)").unwrap().to_string(),
617 std::env::var("PROGRAMFILES(x86)").unwrap()
618 );
619 }
620
621 #[test]
622 fn test_global_variables() {
623 let v = Variables::env();
624 let mut p = PowerShellSession::new().with_variables(v);
625
626 p.parse_input(r#" $global:var_int = 5 "#).unwrap();
627 p.parse_input(r#" $global:var_string = "global";$script:var_string = "script";$local:var_string = "local" "#).unwrap();
628
629 assert_eq!(
630 p.parse_input(r#" $var_int "#).unwrap().result(),
631 PsValue::Int(5)
632 );
633 assert_eq!(
634 p.parse_input(r#" $var_string "#).unwrap().result(),
635 PsValue::String("global".into())
636 );
637
638 let global_variables = p.session_variables();
639 assert_eq!(global_variables.get("var_int").unwrap(), &PsValue::Int(5));
640 assert_eq!(
641 global_variables.get("var_string").unwrap(),
642 &PsValue::String("global".into())
643 );
644 }
645
646 #[test]
647 fn test_script_variables() {
648 let v = Variables::env();
649 let mut p = PowerShellSession::new().with_variables(v);
650
651 let script_res = p
652 .parse_input(r#" $script:var_int = 5;$var_string = "assdfa" "#)
653 .unwrap();
654 let script_variables = script_res.script_variables();
655 assert_eq!(script_variables.get("var_int"), Some(&PsValue::Int(5)));
656 assert_eq!(
657 script_variables.get("var_string"),
658 Some(&PsValue::String("assdfa".into()))
659 );
660 }
661
662 #[test]
663 fn test_env_special_cases() {
664 let v = Variables::env();
665 let mut p = PowerShellSession::new().with_variables(v);
666 p.safe_eval(r#" $global:program = $env:programfiles + "\program" "#)
667 .unwrap();
668 assert_eq!(
669 p.safe_eval(r#" $global:program "#).unwrap().as_str(),
670 format!("{}\\program", std::env::var("PROGRAMFILES").unwrap())
671 );
672 assert_eq!(
673 p.safe_eval(r#" $program "#).unwrap().as_str(),
674 format!("{}\\program", std::env::var("PROGRAMFILES").unwrap())
675 );
676
677 assert_eq!(
678 p.safe_eval(r#" ${Env:ProgramFiles(x86):adsf} = 5;${Env:ProgramFiles(x86):adsf} "#)
679 .unwrap()
680 .as_str(),
681 5.to_string()
682 );
683 assert_eq!(
684 p.safe_eval(r#" ${Env:ProgramFiles(x86)} "#)
685 .unwrap()
686 .as_str(),
687 std::env::var("ProgramFiles(x86)").unwrap()
688 );
689 }
690
691 #[test]
692 fn special_last_error() {
693 let input = r#"3+"01234 ?";$a=5;$a;$?"#;
694
695 let mut p = PowerShellSession::new();
696 assert_eq!(p.safe_eval(input).unwrap().as_str(), "True");
697
698 let input = r#"3+"01234 ?";$?"#;
699 assert_eq!(p.safe_eval(input).unwrap().as_str(), "False");
700 }
701
702 #[test]
703 fn test_from_ini() {
704 let input = r#"[global]
705name = radek
706age = 30
707is_admin = true
708height = 5.9
709empty_value =
710
711[script]
712local_var = "local_value"
713 "#;
714 let mut variables = Variables::new().values_persist();
715 variables.load_from_string(input).unwrap();
716 let mut p = PowerShellSession::new().with_variables(variables);
717
718 assert_eq!(
719 p.parse_input(r#" $global:name "#).unwrap().result(),
720 PsValue::String("radek".into())
721 );
722 assert_eq!(
723 p.parse_input(r#" $global:age "#).unwrap().result(),
724 PsValue::Int(30)
725 );
726 assert_eq!(p.safe_eval(r#" $false "#).unwrap().as_str(), "False");
727 assert_eq!(p.safe_eval(r#" $null "#).unwrap().as_str(), "");
728 assert_eq!(
729 p.safe_eval(r#" $script:local_var "#).unwrap().as_str(),
730 "\"local_value\""
731 );
732 assert_eq!(
733 p.safe_eval(r#" $local:local_var "#).unwrap().as_str(),
734 "\"local_value\""
735 );
736 }
737
738 #[test]
739 fn test_from_ini_string() {
740 let input = r#"[global]
741name = radek
742age = 30
743is_admin = true
744height = 5.9
745empty_value =
746
747[script]
748local_var = "local_value"
749 "#;
750
751 let variables = Variables::from_ini_string(input).unwrap().values_persist();
752 let mut p = PowerShellSession::new().with_variables(variables);
753 assert_eq!(
754 p.parse_input(r#" $global:name "#).unwrap().result(),
755 PsValue::String("radek".into())
756 );
757 assert_eq!(
758 p.parse_input(r#" $global:age "#).unwrap().result(),
759 PsValue::Int(30)
760 );
761 assert_eq!(p.safe_eval(r#" $false "#).unwrap().as_str(), "False");
762 assert_eq!(p.safe_eval(r#" $null "#).unwrap().as_str(), "");
763 assert_eq!(
764 p.safe_eval(r#" $script:local_var "#).unwrap().as_str(),
765 "\"local_value\""
766 );
767 assert_eq!(
768 p.safe_eval(r#" $local_var "#).unwrap().as_str(),
769 "\"local_value\""
770 );
771 assert_eq!(
772 p.safe_eval(r#" $local:local_var "#).unwrap().as_str(),
773 "\"local_value\""
774 );
775 }
776}