use serde_json::json;
use crate::cdp::attach::CdpAttach;
use crate::cdp::options::ChromiumContextOverride;
use crate::cdp::tab::ChromiumTab;
use crate::cdp::types::CookieParam;
use crate::{Error, Result};
#[derive(Debug, Clone, Default)]
pub struct ContextSpec {
pub name: Option<String>,
pub proxy: Option<String>,
pub proxy_bypass: Option<String>,
pub user_agent: Option<String>,
pub locale: Option<String>,
pub timezone: Option<String>,
pub extra_headers: Vec<(String, String)>,
pub permissions: Vec<(String, Vec<String>)>,
pub viewport: Option<(u32, u32)>,
}
impl ContextSpec {
pub fn new() -> Self {
Self::default()
}
pub fn name(mut self, name: impl Into<String>) -> Self {
self.name = Some(name.into());
self
}
pub fn proxy(mut self, p: impl Into<String>) -> Self {
self.proxy = Some(p.into());
self
}
pub fn proxy_bypass(mut self, b: impl Into<String>) -> Self {
self.proxy_bypass = Some(b.into());
self
}
pub fn user_agent(mut self, ua: impl Into<String>) -> Self {
self.user_agent = Some(ua.into());
self
}
pub fn locale(mut self, l: impl Into<String>) -> Self {
self.locale = Some(l.into());
self
}
pub fn timezone(mut self, tz: impl Into<String>) -> Self {
self.timezone = Some(tz.into());
self
}
pub fn extra_header(mut self, name: impl Into<String>, value: impl Into<String>) -> Self {
self.extra_headers.push((name.into(), value.into()));
self
}
pub fn extra_headers(mut self, headers: Vec<(String, String)>) -> Self {
self.extra_headers = headers;
self
}
pub fn grant(mut self, origin: impl Into<String>, permissions: &[&str]) -> Self {
self.permissions.push((
origin.into(),
permissions.iter().map(|s| (*s).to_string()).collect(),
));
self
}
pub fn viewport(mut self, width: u32, height: u32) -> Self {
self.viewport = Some((width, height));
self
}
pub fn from_override(ov: &ChromiumContextOverride) -> Self {
Self {
name: None,
proxy: ov.proxy.clone(),
proxy_bypass: ov.proxy_bypass.clone(),
user_agent: ov.user_agent.clone(),
locale: ov.locale.clone(),
timezone: ov.timezone.clone(),
extra_headers: Vec::new(),
permissions: Vec::new(),
viewport: None,
}
}
}
impl From<&ChromiumContextOverride> for ContextSpec {
fn from(ov: &ChromiumContextOverride) -> Self {
Self::from_override(ov)
}
}
#[derive(Clone)]
pub struct ChromiumBrowserContext {
attach: CdpAttach,
id: String,
spec: ContextSpec,
}
impl ChromiumBrowserContext {
pub(crate) async fn create(attach: CdpAttach, spec: ContextSpec) -> Result<Self> {
if let Some(name) = spec.name.as_deref() {
if let Some(id) = attach.lookup_named(name) {
return Ok(Self { attach, id, spec });
}
}
let mut params = json!({});
if let Some(proxy) = &spec.proxy {
params["proxyServer"] = json!(proxy);
if let Some(b) = &spec.proxy_bypass {
params["proxyBypassList"] = json!(b);
}
}
let r = attach
.conn
.send("Target.createBrowserContext", params, None)
.await?;
let id = r["browserContextId"]
.as_str()
.ok_or_else(|| Error::msg("CDP: createBrowserContext 未返回 browserContextId"))?
.to_string();
if let Some(name) = spec.name.as_deref() {
attach.register_named(name, &id);
}
let ctx = Self { attach, id, spec };
ctx.apply_permissions().await;
Ok(ctx)
}
pub fn id(&self) -> &str {
&self.id
}
pub fn name(&self) -> Option<&str> {
self.spec.name.as_deref()
}
pub fn spec(&self) -> &ContextSpec {
&self.spec
}
pub async fn new_tab(&self, url: Option<&str>) -> Result<ChromiumTab> {
let mut tparams = json!({ "url": url.unwrap_or("about:blank") });
tparams["browserContextId"] = json!(self.id);
let r = self
.attach
.conn
.send("Target.createTarget", tparams, None)
.await?;
let target_id = r["targetId"]
.as_str()
.ok_or_else(|| Error::msg("CDP: 创建标签无 targetId"))?
.to_string();
let tab = self
.attach
.attach_in_context(target_id, Some(self.id.clone()), false)
.await?;
self.apply_session(&tab).await;
Ok(tab)
}
pub async fn set_cookies(&self, cookies: Vec<CookieParam>) -> Result<()> {
let arr: Vec<serde_json::Value> = cookies
.iter()
.map(|c| {
let mut o = json!({ "name": c.name, "value": c.value });
if let Some(u) = &c.url {
o["url"] = json!(u);
}
if let Some(d) = &c.domain {
o["domain"] = json!(d);
}
if let Some(p) = &c.path {
o["path"] = json!(p);
}
if let Some(v) = c.secure {
o["secure"] = json!(v);
}
if let Some(v) = c.http_only {
o["httpOnly"] = json!(v);
}
if let Some(v) = c.expires {
o["expires"] = json!(v);
}
o
})
.collect();
self.attach
.conn
.send(
"Storage.setCookies",
json!({ "cookies": arr, "browserContextId": self.id }),
None,
)
.await?;
Ok(())
}
pub async fn cookies(&self) -> Result<Vec<crate::cdp::types::Cookie>> {
let r = self
.attach
.conn
.send(
"Storage.getCookies",
json!({ "browserContextId": self.id }),
None,
)
.await?;
let s = |c: &serde_json::Value, k: &str| {
c.get(k)
.and_then(serde_json::Value::as_str)
.unwrap_or("")
.to_string()
};
Ok(r["cookies"]
.as_array()
.cloned()
.unwrap_or_default()
.iter()
.map(|c| crate::cdp::types::Cookie {
name: s(c, "name"),
value: s(c, "value"),
domain: s(c, "domain"),
path: s(c, "path"),
expires: c
.get("expires")
.and_then(serde_json::Value::as_f64)
.unwrap_or(-1.0),
http_only: c
.get("httpOnly")
.and_then(serde_json::Value::as_bool)
.unwrap_or(false),
secure: c
.get("secure")
.and_then(serde_json::Value::as_bool)
.unwrap_or(false),
})
.collect())
}
pub async fn close(self) -> Result<()> {
if let Some(name) = self.spec.name.as_deref() {
self.attach.unregister_named(name);
}
self.attach
.conn
.send(
"Target.disposeBrowserContext",
json!({ "browserContextId": self.id }),
None,
)
.await?;
Ok(())
}
async fn apply_permissions(&self) {
for (origin, perms) in &self.spec.permissions {
let mut p = json!({ "permissions": perms, "browserContextId": self.id });
if !origin.is_empty() {
p["origin"] = json!(origin);
}
let _ = self
.attach
.conn
.send("Browser.grantPermissions", p, None)
.await;
}
}
async fn apply_session(&self, tab: &ChromiumTab) {
let spec = ChromiumContextOverride {
proxy: None,
proxy_bypass: None,
user_agent: self.spec.user_agent.clone(),
locale: self.spec.locale.clone(),
timezone: self.spec.timezone.clone(),
};
spec.apply_emulation(tab).await;
if !self.spec.extra_headers.is_empty() {
let mut obj = serde_json::Map::new();
for (k, v) in &self.spec.extra_headers {
obj.insert(k.clone(), json!(v));
}
let _ = tab.core.send("Network.enable", json!({})).await;
let _ = tab
.core
.send(
"Network.setExtraHTTPHeaders",
json!({ "headers": serde_json::Value::Object(obj) }),
)
.await;
}
if let Some((w, h)) = self.spec.viewport {
let _ = tab
.core
.send(
"Emulation.setDeviceMetricsOverride",
json!({
"width": w,
"height": h,
"deviceScaleFactor": 1,
"mobile": false
}),
)
.await;
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn spec_builder_chains() {
let s = ContextSpec::new()
.name("account_001")
.proxy("http://127.0.0.1:8080")
.user_agent("UA/1")
.locale("zh-CN")
.timezone("Asia/Shanghai")
.extra_header("X-Test", "1")
.grant("https://example.com", &["geolocation"])
.viewport(1280, 800);
assert_eq!(s.name.as_deref(), Some("account_001"));
assert_eq!(s.proxy.as_deref(), Some("http://127.0.0.1:8080"));
assert_eq!(s.extra_headers.len(), 1);
assert_eq!(s.permissions.len(), 1);
assert_eq!(s.viewport, Some((1280, 800)));
}
}