1use 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 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 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 {
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}