1use crate::element::Element;
2use crate::transport::{Transport, TransportResponse, next_id};
3use crate::types::{CaptureOptions, ImageFormat, Viewport};
4use crate::utils::{self, send_and_get_msg};
5use anyhow::{Context, Result, anyhow};
6use serde_json::{Value, json};
7use std::sync::Arc;
8use std::time::Duration;
9use tokio::time;
10
11pub struct Tab {
13 pub(crate) transport: Arc<Transport>,
14 pub(crate) session_id: String,
15 pub(crate) target_id: String,
16}
17
18impl Tab {
19 pub(crate) async fn new(transport: Arc<Transport>) -> Result<Self> {
20 let TransportResponse::Response(res_create) = transport
21 .send(json!({ "id": next_id(), "method": "Target.createTarget", "params": { "url": "about:blank" } }))
22 .await? else { return Err(anyhow!("Invalid response type")); };
23
24 let target_id = res_create.result["targetId"]
25 .as_str()
26 .context("No targetId")?
27 .to_string();
28
29 let TransportResponse::Response(res_attach) = transport
30 .send(json!({ "id": next_id(), "method": "Target.attachToTarget", "params": { "targetId": target_id } }))
31 .await? else { return Err(anyhow!("Invalid response type")); };
32
33 let session_id = res_attach.result["sessionId"]
34 .as_str()
35 .context("No sessionId")?
36 .to_string();
37
38 Ok(Self {
39 transport,
40 session_id,
41 target_id,
42 })
43 }
44
45 pub(crate) async fn send_cmd(&self, method: &str, params: serde_json::Value) -> Result<Value> {
46 let msg_id = next_id();
47 let msg = json!({
48 "id": msg_id,
49 "method": method,
50 "params": params
51 })
52 .to_string();
53 let res = send_and_get_msg(self.transport.clone(), msg_id, &self.session_id, msg).await?;
54 utils::serde_msg(&res)
55 }
56
57 pub async fn set_viewport(&self, viewport: &Viewport) -> Result<&Self> {
58 let screen_orientation = if viewport.is_landscape {
59 json!({"type": "landscapePrimary", "angle": 90})
60 } else {
61 json!({"type": "portraitPrimary", "angle": 0})
62 };
63
64 self.send_cmd(
65 "Emulation.setDeviceMetricsOverride",
66 json!({
67 "width": viewport.width,
68 "height": viewport.height,
69 "deviceScaleFactor": viewport.device_scale_factor,
70 "mobile": viewport.is_mobile,
71 "screenOrientation": screen_orientation
72 }),
73 )
74 .await?;
75
76 if viewport.has_touch {
77 self.send_cmd(
78 "Emulation.setTouchEmulationEnabled",
79 json!({
80 "enabled": true,
81 "maxTouchPoints": 5
82 }),
83 )
84 .await?;
85 }
86
87 Ok(self)
88 }
89
90 pub async fn clear_viewport(&self) -> Result<&Self> {
91 self.send_cmd("Emulation.clearDeviceMetricsOverride", json!({}))
92 .await?;
93 Ok(self)
94 }
95
96 pub async fn set_content(&self, content: &str) -> Result<&Self> {
97 self.send_cmd("Page.enable", json!({})).await?;
98
99 let event_rx = self
101 .transport
102 .listen_for_event(&self.session_id, "Page.loadEventFired")
103 .await?;
104
105 let js_write = format!(
106 r#"document.open(); document.write({}); document.close();"#,
107 serde_json::to_string(content)?
108 );
109
110 self.send_cmd(
111 "Runtime.evaluate",
112 json!({
113 "expression": js_write,
114 "awaitPromise": true
115 }),
116 )
117 .await?;
118
119 time::timeout(Duration::from_secs(30), event_rx)
120 .await
121 .map_err(|_| anyhow!("Timeout waiting for event Page.loadEventFired"))?
122 .map_err(|_| anyhow!("Event channel closed"))?;
123
124 Ok(self)
125 }
126
127 pub async fn evaluate(&self, expression: &str) -> Result<Value> {
128 let result = self
129 .send_cmd(
130 "Runtime.evaluate",
131 json!({
132 "expression": expression,
133 "returnByValue": true,
134 "awaitPromise": true
135 }),
136 )
137 .await?;
138 Ok(result["result"]["result"]["value"].clone())
139 }
140
141 pub async fn evaluate_as_string(&self, expression: &str) -> Result<String> {
142 let value = self.evaluate(expression).await?;
143 value
144 .as_str()
145 .map(|s| s.to_string())
146 .or_else(|| Some(value.to_string()))
147 .context("Failed to convert result to string")
148 }
149
150 pub async fn find_element(&self, selector: &str) -> Result<Element<'_>> {
151 let msg_id = next_id();
152 let msg_doc =
153 json!({ "id": msg_id, "method": "DOM.getDocument", "params": {} }).to_string();
154 let res_doc =
155 send_and_get_msg(self.transport.clone(), msg_id, &self.session_id, msg_doc).await?;
156 let data_doc = utils::serde_msg(&res_doc)?;
157 let root_node_id = data_doc["result"]["root"]["nodeId"]
158 .as_u64()
159 .context("No root node")?;
160
161 let msg_id = next_id();
162 let msg_sel = json!({
163 "id": msg_id,
164 "method": "DOM.querySelector",
165 "params": { "nodeId": root_node_id, "selector": selector }
166 })
167 .to_string();
168 let res_sel =
169 send_and_get_msg(self.transport.clone(), msg_id, &self.session_id, msg_sel).await?;
170 let data_sel = utils::serde_msg(&res_sel)?;
171 let node_id = data_sel["result"]["nodeId"]
172 .as_u64()
173 .context("Element not found")?;
174
175 Element::new(self, node_id).await
176 }
177
178 pub async fn wait_for_selector(&self, selector: &str, timeout_ms: u64) -> Result<Element<'_>> {
179 let start = std::time::Instant::now();
180 let timeout = std::time::Duration::from_millis(timeout_ms);
181
182 loop {
183 match self.find_element(selector).await {
184 Ok(element) => return Ok(element),
185 Err(_) if start.elapsed() < timeout => {
186 tokio::time::sleep(std::time::Duration::from_millis(100)).await;
187 }
188 Err(e) => return Err(e),
189 }
190 }
191 }
192
193 pub async fn screenshot(&self, opts: CaptureOptions) -> Result<String> {
194 if let Some(ref viewport) = opts.viewport {
195 self.set_viewport(viewport).await?;
196 }
197
198 let mut params = json!({
199 "format": opts.format.as_str(),
200 "fromSurface": true,
201 "captureBeyondViewport": opts.full_page,
202 });
203
204 if matches!(opts.format, ImageFormat::Jpeg | ImageFormat::WebP) {
205 params["quality"] = json!(opts.quality.unwrap_or(90));
206 }
207
208 if let Some(ref clip) = opts.clip {
209 params["clip"] = json!({
210 "x": clip.x,
211 "y": clip.y,
212 "width": clip.width,
213 "height": clip.height,
214 "scale": clip.scale
215 });
216 }
217
218 if opts.omit_background && matches!(opts.format, ImageFormat::Png) {
219 self.send_cmd(
220 "Emulation.setDefaultBackgroundColorOverride",
221 json!({ "color": { "r": 0, "g": 0, "b": 0, "a": 0 } }),
222 )
223 .await?;
224 }
225
226 self.activate().await?;
227
228 let result = self.send_cmd("Page.captureScreenshot", params).await?;
229
230 if opts.omit_background && matches!(opts.format, ImageFormat::Png) {
231 let _ = self
232 .send_cmd("Emulation.setDefaultBackgroundColorOverride", json!({}))
233 .await;
234 }
235
236 result["result"]["data"]
237 .as_str()
238 .map(|s| s.to_string())
239 .context("No image data received")
240 }
241
242 pub async fn activate(&self) -> Result<&Self> {
243 let msg_id = next_id();
244 let msg = json!({ "id": msg_id, "method": "Target.activateTarget", "params": { "targetId": self.target_id } }).to_string();
245 send_and_get_msg(self.transport.clone(), msg_id, &self.session_id, msg).await?;
246 Ok(self)
247 }
248
249 pub async fn goto(&self, url: &str) -> Result<&Self> {
250 self.send_cmd("Page.enable", json!({})).await?;
251
252 let event_rx = self
253 .transport
254 .listen_for_event(&self.session_id, "Page.loadEventFired")
255 .await?;
256
257 let msg_id = next_id();
258 let msg = json!({ "id": msg_id, "method": "Page.navigate", "params": { "url": url } })
259 .to_string();
260 send_and_get_msg(self.transport.clone(), msg_id, &self.session_id, msg).await?;
261
262 time::timeout(Duration::from_secs(30), event_rx)
263 .await
264 .map_err(|_| anyhow!("Timeout waiting for event Page.loadEventFired"))?
265 .map_err(|_| anyhow!("Event channel closed"))?;
266
267 Ok(self)
268 }
269
270 pub async fn goto_no_wait(&self, url: &str) -> Result<&Self> {
271 let msg_id = next_id();
272 let msg = json!({ "id": msg_id, "method": "Page.navigate", "params": { "url": url } })
273 .to_string();
274 send_and_get_msg(self.transport.clone(), msg_id, &self.session_id, msg).await?;
275 Ok(self)
276 }
277
278 pub async fn reload(&self) -> Result<&Self> {
279 self.send_cmd("Page.enable", json!({})).await?;
280
281 let event_rx = self
282 .transport
283 .listen_for_event(&self.session_id, "Page.loadEventFired")
284 .await?;
285
286 self.send_cmd("Page.reload", json!({})).await?;
287
288 time::timeout(Duration::from_secs(30), event_rx)
289 .await
290 .map_err(|_| anyhow!("Timeout waiting for event Page.loadEventFired"))?
291 .map_err(|_| anyhow!("Event channel closed"))?;
292
293 Ok(self)
294 }
295
296 pub async fn url(&self) -> Result<String> {
297 self.evaluate_as_string("window.location.href").await
298 }
299
300 pub async fn title(&self) -> Result<String> {
301 self.evaluate_as_string("document.title").await
302 }
303
304 pub fn session_id(&self) -> &str {
305 &self.session_id
306 }
307
308 pub fn target_id(&self) -> &str {
309 &self.target_id
310 }
311
312 pub async fn close(&self) -> Result<()> {
313 let msg_id = next_id();
314 let msg = json!({ "id": msg_id, "method": "Target.closeTarget", "params": { "targetId": self.target_id } }).to_string();
315 send_and_get_msg(self.transport.clone(), msg_id, &self.session_id, msg).await?;
316 Ok(())
317 }
318}