use serde::{Deserialize, Serialize};
use crate::browser::Cookie;
use crate::browser::tab::{CookieParam, Tab};
use crate::Result;
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct StorageState {
#[serde(default)]
pub cookies: Vec<CookieParam>,
#[serde(default)]
pub origins: Vec<OriginStorage>,
}
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct OriginStorage {
pub origin: String,
#[serde(default)]
pub local_storage: Vec<(String, String)>,
#[serde(default)]
pub session_storage: Vec<(String, String)>,
}
impl Tab {
pub async fn storage_state(&self) -> Result<StorageState> {
let cookies = self
.cookies()
.await?
.into_iter()
.map(cookie_to_param)
.collect();
let origin = self
.run_js("location.origin")
.await?
.as_str()
.unwrap_or_default()
.to_string();
let local = read_storage(self, "localStorage").await;
let session = read_storage(self, "sessionStorage").await;
let mut origins = Vec::new();
if !origin.is_empty() && origin != "null" && (!local.is_empty() || !session.is_empty()) {
origins.push(OriginStorage {
origin,
local_storage: local,
session_storage: session,
});
}
Ok(StorageState { cookies, origins })
}
pub async fn save_storage_state(&self, path: &str) -> Result<()> {
let st = self.storage_state().await?;
std::fs::write(path, serde_json::to_string_pretty(&st)?)?;
Ok(())
}
pub async fn load_storage_state(&self, path: &str) -> Result<()> {
let s = std::fs::read_to_string(path)?;
let st: StorageState = serde_json::from_str(&s)?;
self.apply_storage_state(&st).await
}
pub async fn apply_storage_state(&self, state: &StorageState) -> Result<()> {
if !state.cookies.is_empty() {
self.set_cookies(state.cookies.clone()).await?;
}
let origin = self
.run_js("location.origin")
.await?
.as_str()
.unwrap_or_default()
.to_string();
if let Some(os) = state.origins.iter().find(|o| o.origin == origin) {
write_storage(self, "localStorage", &os.local_storage).await?;
write_storage(self, "sessionStorage", &os.session_storage).await?;
}
Ok(())
}
}
fn cookie_to_param(c: Cookie) -> CookieParam {
CookieParam {
name: c.name,
value: c.value,
url: None,
domain: Some(c.domain),
path: Some(c.path),
secure: Some(c.secure),
http_only: Some(c.http_only),
expires: if c.expires > 0.0 { Some(c.expires) } else { None },
}
}
async fn read_storage(tab: &Tab, which: &str) -> Vec<(String, String)> {
let js = format!(
"(function(){{ try {{ var s={which}; var o=[]; for(var i=0;i<s.length;i++){{ var k=s.key(i); o.push([k, s.getItem(k)]); }} return JSON.stringify(o); }} catch(e){{ return '[]'; }} }})()"
);
let v = match tab.run_js(&js).await {
Ok(v) => v,
Err(_) => return Vec::new(),
};
let s = v.as_str().unwrap_or("[]");
serde_json::from_str::<Vec<[String; 2]>>(s)
.map(|rows| rows.into_iter().map(|[k, v]| (k, v)).collect())
.unwrap_or_default()
}
async fn write_storage(tab: &Tab, which: &str, items: &[(String, String)]) -> Result<()> {
if items.is_empty() {
return Ok(());
}
let json = serde_json::to_string(items).unwrap_or_else(|_| "[]".into());
let js = format!(
"(function(){{ try {{ var d={json}; for(var i=0;i<d.length;i++){{ {which}.setItem(d[i][0], d[i][1]); }} }} catch(e){{}} return true; }})()"
);
tab.run_js(&js).await?;
Ok(())
}