Skip to main content

strat9_kernel/shell/
scripting.rs

1//! Minimal shell scripting: variable expansion, `if`, `for`, `while`.
2//!
3//! # Supported constructs
4//!
5//! - **Variables**: `set VAR=VALUE`, `$VAR` expansion, `$?` for last exit code.
6//! - **For loops**: `for VAR in A B C ; do COMMAND ; done`
7//! - **While loops**: `while COMMAND ; do BODY ; done`
8//! - **If/then**: `if COMMAND ; then BODY ; fi` or `if COMMAND ; then A ; else B ; fi`
9
10use crate::sync::SpinLock;
11use alloc::{
12    collections::BTreeMap,
13    string::{String, ToString},
14    vec::Vec,
15};
16
17static SHELL_VARS: SpinLock<BTreeMap<String, String>> = SpinLock::new(BTreeMap::new());
18static LAST_EXIT: core::sync::atomic::AtomicI32 = core::sync::atomic::AtomicI32::new(0);
19
20/// Set the exit code of the last executed command.
21pub fn set_last_exit(code: i32) {
22    LAST_EXIT.store(code, core::sync::atomic::Ordering::Relaxed);
23}
24
25/// Get the exit code of the last executed command.
26pub fn last_exit() -> i32 {
27    LAST_EXIT.load(core::sync::atomic::Ordering::Relaxed)
28}
29
30/// Set a shell variable.
31pub fn set_var(key: &str, val: &str) {
32    SHELL_VARS
33        .lock()
34        .insert(String::from(key), String::from(val));
35}
36
37/// Get a shell variable.
38pub fn get_var(key: &str) -> Option<String> {
39    SHELL_VARS.lock().get(key).cloned()
40}
41
42/// Remove a shell variable.
43pub fn unset_var(key: &str) {
44    SHELL_VARS.lock().remove(key);
45}
46
47/// Expand `$VAR` and `$?` references in a string.
48pub fn expand_vars(input: &str) -> String {
49    let mut result = String::new();
50    let bytes = input.as_bytes();
51    let mut i = 0;
52
53    while i < bytes.len() {
54        if bytes[i] == b'$' {
55            i += 1;
56            if i < bytes.len() && bytes[i] == b'?' {
57                result.push_str(&alloc::format!("{}", last_exit()));
58                i += 1;
59            } else {
60                let start = i;
61                while i < bytes.len() && (bytes[i].is_ascii_alphanumeric() || bytes[i] == b'_') {
62                    i += 1;
63                }
64                if i == start {
65                    // Keep literal '$' when no variable name follows.
66                    result.push('$');
67                    continue;
68                }
69                let var_name = core::str::from_utf8(&bytes[start..i]).unwrap_or("");
70                if let Some(val) = get_var(var_name) {
71                    result.push_str(&val);
72                } else if let Some(env_val) = super::commands::util::shell_getenv(var_name) {
73                    result.push_str(&env_val);
74                }
75            }
76        } else {
77            result.push(bytes[i] as char);
78            i += 1;
79        }
80    }
81
82    result
83}
84
85/// Script construct types recognized by the parser.
86pub enum ScriptConstruct {
87    /// Simple command line (possibly with pipes).
88    Simple(String),
89    /// `set VAR=VALUE`
90    SetVar { key: String, val: String },
91    /// `unset VAR`
92    UnsetVar(String),
93    /// `for VAR in ITEMS ; do BODY ; done`
94    ForLoop {
95        var: String,
96        items: Vec<String>,
97        body: Vec<String>,
98    },
99    /// `while COND ; do BODY ; done`
100    WhileLoop { cond: String, body: Vec<String> },
101    /// `if COND ; then THEN_BODY [; else ELSE_BODY] ; fi`
102    IfElse {
103        cond: String,
104        then_body: Vec<String>,
105        else_body: Vec<String>,
106    },
107}
108
109/// Parse a full line into a script construct.
110///
111/// The parser splits on `;` to find keywords. For simple commands
112/// (no scripting keywords), returns [`ScriptConstruct::Simple`].
113pub fn parse_script(line: &str) -> ScriptConstruct {
114    let trimmed = line.trim();
115
116    if trimmed.starts_with("set ") {
117        let rest = &trimmed[4..];
118        if let Some(eq) = rest.find('=') {
119            return ScriptConstruct::SetVar {
120                key: String::from(rest[..eq].trim()),
121                val: String::from(rest[eq + 1..].trim()),
122            };
123        }
124        return ScriptConstruct::Simple(String::from(trimmed));
125    }
126
127    if trimmed.starts_with("unset ") {
128        return ScriptConstruct::UnsetVar(String::from(trimmed[6..].trim()));
129    }
130
131    let parts: Vec<&str> = trimmed.split(';').map(|s| s.trim()).collect();
132
133    if parts.first().map(|p| p.starts_with("for ")) == Some(true) {
134        return parse_for_loop(&parts);
135    }
136    if parts.first().map(|p| p.starts_with("while ")) == Some(true) {
137        return parse_while_loop(&parts);
138    }
139    if parts.first().map(|p| p.starts_with("if ")) == Some(true) {
140        return parse_if_else(&parts);
141    }
142
143    ScriptConstruct::Simple(String::from(trimmed))
144}
145
146/// `for VAR in A B C ; do cmd1 ; cmd2 ; done`
147fn parse_for_loop(parts: &[&str]) -> ScriptConstruct {
148    if parts.is_empty() {
149        return ScriptConstruct::Simple(String::new());
150    }
151    let header = parts[0];
152    let tokens: Vec<&str> = header.split_whitespace().collect();
153    if tokens.len() < 4 || tokens[0] != "for" || tokens[2] != "in" {
154        return ScriptConstruct::Simple(String::from(header));
155    }
156
157    let var = tokens.get(1).unwrap_or(&"_").to_string();
158    let items: Vec<String> = tokens
159        .iter()
160        .skip(3) // skip "for VAR in"
161        .map(|s| String::from(*s))
162        .collect();
163
164    let Some(do_idx) = parts.iter().position(|p| *p == "do") else {
165        return ScriptConstruct::Simple(String::from(header));
166    };
167    let Some(done_idx) = parts.iter().position(|p| *p == "done") else {
168        return ScriptConstruct::Simple(String::from(header));
169    };
170    if done_idx <= do_idx {
171        return ScriptConstruct::Simple(String::from(header));
172    }
173
174    let body: Vec<String> = parts[do_idx + 1..done_idx]
175        .iter()
176        .filter(|s| !s.is_empty())
177        .map(|s| String::from(*s))
178        .collect();
179
180    ScriptConstruct::ForLoop {
181        var: String::from(var),
182        items,
183        body,
184    }
185}
186
187/// `while cond ; do body ; done`
188fn parse_while_loop(parts: &[&str]) -> ScriptConstruct {
189    if parts.is_empty() {
190        return ScriptConstruct::Simple(String::new());
191    }
192    let header = parts[0];
193    let Some(cond) = header.strip_prefix("while ") else {
194        return ScriptConstruct::Simple(String::from(header));
195    };
196    let Some(do_idx) = parts.iter().position(|p| *p == "do") else {
197        return ScriptConstruct::Simple(String::from(header));
198    };
199    let Some(done_idx) = parts.iter().position(|p| *p == "done") else {
200        return ScriptConstruct::Simple(String::from(header));
201    };
202    if done_idx <= do_idx {
203        return ScriptConstruct::Simple(String::from(header));
204    }
205
206    let body: Vec<String> = parts[do_idx + 1..done_idx]
207        .iter()
208        .filter(|s| !s.is_empty())
209        .map(|s| String::from(*s))
210        .collect();
211
212    ScriptConstruct::WhileLoop {
213        cond: String::from(cond),
214        body,
215    }
216}
217
218/// `if cond ; then body ; [else body ;] fi`
219fn parse_if_else(parts: &[&str]) -> ScriptConstruct {
220    if parts.is_empty() {
221        return ScriptConstruct::Simple(String::new());
222    }
223    let header = parts[0];
224    let Some(cond) = header.strip_prefix("if ") else {
225        return ScriptConstruct::Simple(String::from(header));
226    };
227    let Some(then_idx) = parts.iter().position(|p| *p == "then") else {
228        return ScriptConstruct::Simple(String::from(header));
229    };
230    let else_idx = parts.iter().position(|p| *p == "else");
231    let Some(fi_idx) = parts.iter().position(|p| *p == "fi") else {
232        return ScriptConstruct::Simple(String::from(header));
233    };
234    if fi_idx <= then_idx {
235        return ScriptConstruct::Simple(String::from(header));
236    }
237
238    let then_end = else_idx.unwrap_or(fi_idx);
239    let then_body: Vec<String> = parts[then_idx + 1..then_end]
240        .iter()
241        .filter(|s| !s.is_empty())
242        .map(|s| String::from(*s))
243        .collect();
244
245    let else_body: Vec<String> = if let Some(ei) = else_idx {
246        parts[ei + 1..fi_idx]
247            .iter()
248            .filter(|s| !s.is_empty())
249            .map(|s| String::from(*s))
250            .collect()
251    } else {
252        Vec::new()
253    };
254
255    ScriptConstruct::IfElse {
256        cond: String::from(cond),
257        then_body,
258        else_body,
259    }
260}