Skip to main content

browser_control/cli/
storage.rs

1//! `browser-control storage` — local/sessionStorage get/set/list.
2
3use anyhow::{anyhow, bail, Result};
4use clap::Subcommand;
5use serde_json::Value;
6
7use crate::cli::mcp::resolve_browser;
8use crate::session::PageSession;
9
10#[derive(Subcommand, Debug)]
11pub enum StorageCmd {
12    /// Read a single storage entry. With --key-regex, returns the first match.
13    Get {
14        #[arg(env = "BROWSER_CONTROL")]
15        browser: Option<String>,
16        key: Option<String>,
17        #[arg(long)]
18        key_regex: Option<String>,
19        #[arg(long)]
20        target: Option<String>,
21        #[arg(long, default_value = "local")]
22        namespace: String,
23        #[arg(long)]
24        json: bool,
25    },
26    /// Write a single storage entry.
27    Set {
28        #[arg(env = "BROWSER_CONTROL")]
29        browser: Option<String>,
30        key: String,
31        value: String,
32        #[arg(long)]
33        target: Option<String>,
34        #[arg(long, default_value = "local")]
35        namespace: String,
36    },
37    /// List all storage entries (optionally filtered by key regex).
38    List {
39        #[arg(env = "BROWSER_CONTROL")]
40        browser: Option<String>,
41        #[arg(long)]
42        key_regex: Option<String>,
43        #[arg(long)]
44        target: Option<String>,
45        #[arg(long, default_value = "local")]
46        namespace: String,
47        #[arg(long)]
48        json: bool,
49    },
50}
51
52pub async fn run(cmd: StorageCmd) -> Result<()> {
53    match cmd {
54        StorageCmd::Get {
55            browser,
56            key,
57            key_regex,
58            target,
59            namespace,
60            json,
61        } => run_get(browser, key, key_regex, target, namespace, json).await,
62        StorageCmd::Set {
63            browser,
64            key,
65            value,
66            target,
67            namespace,
68        } => run_set(browser, key, value, target, namespace).await,
69        StorageCmd::List {
70            browser,
71            key_regex,
72            target,
73            namespace,
74            json,
75        } => run_list(browser, key_regex, target, namespace, json).await,
76    }
77}
78
79async fn evaluate_in_target(
80    browser: Option<String>,
81    target: Option<String>,
82    expr: &str,
83) -> Result<Value> {
84    let resolved = resolve_browser(browser).await?;
85    let session =
86        PageSession::attach(&resolved.endpoint, resolved.engine, target.as_deref()).await?;
87    let value = session.evaluate(expr, true).await;
88    session.close().await;
89    value
90}
91
92async fn run_get(
93    browser: Option<String>,
94    key: Option<String>,
95    key_regex: Option<String>,
96    target: Option<String>,
97    namespace: String,
98    json: bool,
99) -> Result<()> {
100    let ns = ns_global(&namespace)?;
101    match (key.as_deref(), key_regex.as_deref()) {
102        (Some(_), Some(_)) => bail!("specify either KEY or --key-regex, not both"),
103        (None, None) => bail!("specify a KEY or --key-regex"),
104        (Some(k), None) => {
105            let expr = build_get_expr(ns, k);
106            let value = evaluate_in_target(browser, target, &expr).await?;
107            if value.is_null() {
108                bail!("key not found: {k}");
109            }
110            if json {
111                println!("{}", serde_json::to_string_pretty(&value)?);
112            } else if let Some(s) = value.as_str() {
113                println!("{s}");
114            } else {
115                println!("{value}");
116            }
117            Ok(())
118        }
119        (None, Some(pat)) => {
120            let expr = build_get_by_regex_expr(ns, pat);
121            let value = evaluate_in_target(browser, target, &expr).await?;
122            if value.is_null() {
123                bail!("no key matches regex");
124            }
125            if json {
126                println!("{}", serde_json::to_string_pretty(&value)?);
127            } else {
128                let v = value.get("value").unwrap_or(&Value::Null);
129                if let Some(s) = v.as_str() {
130                    println!("{s}");
131                } else {
132                    println!("{v}");
133                }
134            }
135            Ok(())
136        }
137    }
138}
139
140async fn run_set(
141    browser: Option<String>,
142    key: String,
143    value: String,
144    target: Option<String>,
145    namespace: String,
146) -> Result<()> {
147    let ns = ns_global(&namespace)?;
148    let expr = build_set_expr(ns, &key, &value);
149    evaluate_in_target(browser, target, &expr).await?;
150    Ok(())
151}
152
153async fn run_list(
154    browser: Option<String>,
155    key_regex: Option<String>,
156    target: Option<String>,
157    namespace: String,
158    json: bool,
159) -> Result<()> {
160    let ns = ns_global(&namespace)?;
161    let expr = build_list_expr(ns, key_regex.as_deref());
162    let value = evaluate_in_target(browser, target, &expr).await?;
163    if json {
164        println!("{}", serde_json::to_string_pretty(&value)?);
165        return Ok(());
166    }
167    let arr = value.as_array().ok_or_else(|| anyhow!("expected array"))?;
168    for entry in arr {
169        let k = entry.get("key").and_then(|v| v.as_str()).unwrap_or("");
170        let v_val = entry.get("value").unwrap_or(&Value::Null);
171        let v_str = match v_val {
172            Value::String(s) => {
173                if s.contains('\t') || s.contains('\n') || s.contains('\r') {
174                    serde_json::to_string(s)?
175                } else {
176                    s.clone()
177                }
178            }
179            other => serde_json::to_string(other)?,
180        };
181        println!("{k}\t{v_str}");
182    }
183    Ok(())
184}
185
186pub(crate) fn ns_global(namespace: &str) -> Result<&'static str> {
187    match namespace {
188        "local" => Ok("localStorage"),
189        "session" => Ok("sessionStorage"),
190        other => bail!("invalid namespace `{other}`: expected `local` or `session`"),
191    }
192}
193
194pub(crate) fn build_get_expr(namespace_js: &str, key: &str) -> String {
195    let key_lit = serde_json::to_string(key).expect("string serialization is infallible");
196    format!("JSON.stringify({namespace_js}.getItem({key_lit}))")
197}
198
199fn build_get_by_regex_expr(namespace_js: &str, pattern: &str) -> String {
200    let pat_lit = serde_json::to_string(pattern).expect("string serialization is infallible");
201    format!(
202        "(() => {{ \
203const re = new RegExp({pat_lit}); \
204const k = Object.keys({namespace_js}).find(k => re.test(k)); \
205return k ? {{key: k, value: {namespace_js}.getItem(k)}} : null; \
206}})()"
207    )
208}
209
210pub(crate) fn build_set_expr(namespace_js: &str, key: &str, value: &str) -> String {
211    let key_lit = serde_json::to_string(key).expect("string serialization is infallible");
212    let val_lit = serde_json::to_string(value).expect("string serialization is infallible");
213    format!("{namespace_js}.setItem({key_lit}, {val_lit})")
214}
215
216fn build_list_expr(namespace_js: &str, pattern: Option<&str>) -> String {
217    let re_expr = match pattern {
218        Some(p) => {
219            let pat_lit = serde_json::to_string(p).expect("string serialization is infallible");
220            format!("new RegExp({pat_lit})")
221        }
222        None => "null".to_string(),
223    };
224    format!(
225        "(() => {{ \
226const ns = {namespace_js}; \
227const re = {re_expr}; \
228const out = []; \
229for (let i = 0; i < ns.length; i++) {{ \
230const k = ns.key(i); \
231if (!re || re.test(k)) out.push({{key: k, value: ns.getItem(k)}}); \
232}} \
233return out; \
234}})()"
235    )
236}
237
238#[cfg(test)]
239mod tests {
240    use super::*;
241
242    #[test]
243    fn ns_global_maps_known() {
244        assert_eq!(ns_global("local").unwrap(), "localStorage");
245        assert_eq!(ns_global("session").unwrap(), "sessionStorage");
246    }
247
248    #[test]
249    fn ns_global_rejects_unknown() {
250        let err = ns_global("cookies").unwrap_err().to_string();
251        assert!(err.contains("invalid namespace"), "got: {err}");
252        assert!(err.contains("cookies"));
253    }
254
255    #[test]
256    fn build_get_expr_escapes_single_quote() {
257        let expr = build_get_expr("localStorage", "it's");
258        assert_eq!(expr, "JSON.stringify(localStorage.getItem(\"it's\"))");
259    }
260
261    #[test]
262    fn build_get_expr_escapes_quote_and_backslash() {
263        let expr = build_get_expr("sessionStorage", "a\"b\\c");
264        assert_eq!(
265            expr,
266            "JSON.stringify(sessionStorage.getItem(\"a\\\"b\\\\c\"))"
267        );
268    }
269
270    #[test]
271    fn build_set_expr_escapes_both() {
272        let expr = build_set_expr("localStorage", "k\"1", "v\\n");
273        assert_eq!(expr, "localStorage.setItem(\"k\\\"1\", \"v\\\\n\")");
274    }
275
276    #[test]
277    fn build_get_by_regex_expr_escapes_quotes() {
278        let expr = build_get_by_regex_expr("localStorage", "^foo\".*$");
279        assert!(
280            expr.contains("new RegExp(\"^foo\\\".*$\")"),
281            "expr was: {expr}"
282        );
283        assert!(expr.contains("Object.keys(localStorage)"));
284    }
285
286    #[test]
287    fn build_list_expr_none_uses_null_regex() {
288        let expr = build_list_expr("localStorage", None);
289        assert!(expr.contains("const re = null;"), "expr: {expr}");
290        assert!(expr.contains("const ns = localStorage;"));
291    }
292
293    #[test]
294    fn build_list_expr_some_escapes_pattern() {
295        let expr = build_list_expr("sessionStorage", Some("a\"b"));
296        assert!(expr.contains("new RegExp(\"a\\\"b\")"), "expr: {expr}");
297        assert!(expr.contains("const ns = sessionStorage;"));
298    }
299}