viewpoint_core/page/content/
mod.rs

1//! Content manipulation functionality.
2//!
3//! This module provides methods for manipulating page HTML content,
4//! injecting scripts and styles.
5
6use std::time::Duration;
7
8use tracing::{debug, info, instrument};
9use viewpoint_cdp::protocol::page::SetDocumentContentParams;
10use viewpoint_cdp::protocol::runtime::EvaluateParams;
11
12use crate::error::PageError;
13use crate::wait::DocumentLoadState;
14
15use super::Page;
16
17/// Script type for injection.
18#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
19pub enum ScriptType {
20    /// Regular script (default).
21    #[default]
22    Script,
23    /// ES6 module.
24    Module,
25}
26
27/// Builder for injecting script tags.
28#[derive(Debug)]
29pub struct ScriptTagBuilder<'a> {
30    page: &'a Page,
31    url: Option<String>,
32    content: Option<String>,
33    script_type: ScriptType,
34}
35
36impl<'a> ScriptTagBuilder<'a> {
37    /// Create a new script tag builder.
38    pub(crate) fn new(page: &'a Page) -> Self {
39        Self {
40            page,
41            url: None,
42            content: None,
43            script_type: ScriptType::default(),
44        }
45    }
46
47    /// Set the script URL.
48    #[must_use]
49    pub fn url(mut self, url: impl Into<String>) -> Self {
50        self.url = Some(url.into());
51        self
52    }
53
54    /// Set the script content.
55    #[must_use]
56    pub fn content(mut self, content: impl Into<String>) -> Self {
57        self.content = Some(content.into());
58        self
59    }
60
61    /// Set the script type.
62    #[must_use]
63    pub fn script_type(mut self, script_type: ScriptType) -> Self {
64        self.script_type = script_type;
65        self
66    }
67
68    /// Inject the script tag into the page.
69    ///
70    /// # Errors
71    ///
72    /// Returns an error if the injection fails.
73    #[instrument(level = "debug", skip(self), fields(has_url = self.url.is_some(), has_content = self.content.is_some()))]
74    pub async fn inject(self) -> Result<(), PageError> {
75        if self.page.is_closed() {
76            return Err(PageError::Closed);
77        }
78
79        let script_js = if let Some(url) = self.url {
80            format!(
81                r"
82                new Promise((resolve, reject) => {{
83                    const script = document.createElement('script');
84                    script.src = '{}';
85                    script.setAttribute('type', '{}');
86                    script.onload = resolve;
87                    script.onerror = reject;
88                    document.head.appendChild(script);
89                }})
90                ",
91                url.replace('\'', "\\'"),
92                if self.script_type == ScriptType::Module {
93                    "module"
94                } else {
95                    "text/javascript"
96                }
97            )
98        } else if let Some(content) = self.content {
99            format!(
100                r"
101                (() => {{
102                    const script = document.createElement('script');
103                    script.textContent = {};
104                    {}
105                    document.head.appendChild(script);
106                }})()
107                ",
108                serde_json::to_string(&content).unwrap_or_default(),
109                if self.script_type == ScriptType::Module {
110                    "script.type = 'module';"
111                } else {
112                    ""
113                }
114            )
115        } else {
116            return Err(PageError::EvaluationFailed(
117                "Either url or content must be provided".to_string(),
118            ));
119        };
120
121        debug!("Injecting script tag");
122
123        self.page
124            .connection()
125            .send_command::<_, serde_json::Value>(
126                "Runtime.evaluate",
127                Some(EvaluateParams {
128                    expression: script_js,
129                    object_group: None,
130                    include_command_line_api: None,
131                    silent: Some(false),
132                    context_id: None,
133                    return_by_value: Some(true),
134                    await_promise: Some(true),
135                }),
136                Some(self.page.session_id()),
137            )
138            .await?;
139
140        Ok(())
141    }
142}
143
144/// Builder for injecting style tags.
145#[derive(Debug)]
146pub struct StyleTagBuilder<'a> {
147    page: &'a Page,
148    url: Option<String>,
149    content: Option<String>,
150}
151
152impl<'a> StyleTagBuilder<'a> {
153    /// Create a new style tag builder.
154    pub(crate) fn new(page: &'a Page) -> Self {
155        Self {
156            page,
157            url: None,
158            content: None,
159        }
160    }
161
162    /// Set the stylesheet URL.
163    #[must_use]
164    pub fn url(mut self, url: impl Into<String>) -> Self {
165        self.url = Some(url.into());
166        self
167    }
168
169    /// Set the CSS content.
170    #[must_use]
171    pub fn content(mut self, content: impl Into<String>) -> Self {
172        self.content = Some(content.into());
173        self
174    }
175
176    /// Inject the style into the page.
177    ///
178    /// # Errors
179    ///
180    /// Returns an error if the injection fails.
181    #[instrument(level = "debug", skip(self), fields(has_url = self.url.is_some(), has_content = self.content.is_some()))]
182    pub async fn inject(self) -> Result<(), PageError> {
183        if self.page.is_closed() {
184            return Err(PageError::Closed);
185        }
186
187        let style_js = if let Some(url) = self.url {
188            format!(
189                r"
190                new Promise((resolve, reject) => {{
191                    const link = document.createElement('link');
192                    link.rel = 'stylesheet';
193                    link.href = '{}';
194                    link.onload = resolve;
195                    link.onerror = reject;
196                    document.head.appendChild(link);
197                }})
198                ",
199                url.replace('\'', "\\'")
200            )
201        } else if let Some(content) = self.content {
202            format!(
203                r"
204                (() => {{
205                    const style = document.createElement('style');
206                    style.textContent = {};
207                    document.head.appendChild(style);
208                }})()
209                ",
210                serde_json::to_string(&content).unwrap_or_default()
211            )
212        } else {
213            return Err(PageError::EvaluationFailed(
214                "Either url or content must be provided".to_string(),
215            ));
216        };
217
218        debug!("Injecting style tag");
219
220        self.page
221            .connection()
222            .send_command::<_, serde_json::Value>(
223                "Runtime.evaluate",
224                Some(EvaluateParams {
225                    expression: style_js,
226                    object_group: None,
227                    include_command_line_api: None,
228                    silent: Some(false),
229                    context_id: None,
230                    return_by_value: Some(true),
231                    await_promise: Some(true),
232                }),
233                Some(self.page.session_id()),
234            )
235            .await?;
236
237        Ok(())
238    }
239}
240
241/// Builder for setting page content.
242#[derive(Debug)]
243pub struct SetContentBuilder<'a> {
244    page: &'a Page,
245    html: String,
246    wait_until: DocumentLoadState,
247    timeout: Duration,
248}
249
250impl<'a> SetContentBuilder<'a> {
251    /// Create a new set content builder.
252    pub(crate) fn new(page: &'a Page, html: String) -> Self {
253        Self {
254            page,
255            html,
256            wait_until: DocumentLoadState::Load,
257            timeout: Duration::from_secs(30),
258        }
259    }
260
261    /// Set the load state to wait for.
262    #[must_use]
263    pub fn wait_until(mut self, state: DocumentLoadState) -> Self {
264        self.wait_until = state;
265        self
266    }
267
268    /// Set the timeout.
269    #[must_use]
270    pub fn timeout(mut self, timeout: Duration) -> Self {
271        self.timeout = timeout;
272        self
273    }
274
275    /// Set the page content.
276    ///
277    /// # Errors
278    ///
279    /// Returns an error if setting content fails.
280    #[instrument(level = "info", skip(self), fields(html_len = self.html.len(), wait_until = ?self.wait_until))]
281    pub async fn set(self) -> Result<(), PageError> {
282        if self.page.is_closed() {
283            return Err(PageError::Closed);
284        }
285
286        info!("Setting page content");
287
288        // Use Page.setDocumentContent
289        self.page
290            .connection()
291            .send_command::<_, serde_json::Value>(
292                "Page.setDocumentContent",
293                Some(SetDocumentContentParams {
294                    frame_id: self.page.frame_id().to_string(),
295                    html: self.html,
296                }),
297                Some(self.page.session_id()),
298            )
299            .await?;
300
301        // Wait for the specified load state
302        // For setDocumentContent, the content is set synchronously,
303        // but we may need to wait for any scripts/resources to load
304        if self.wait_until != DocumentLoadState::Commit {
305            // Small delay to allow the document to settle
306            tokio::time::sleep(Duration::from_millis(50)).await;
307        }
308
309        info!("Page content set");
310        Ok(())
311    }
312}
313
314impl Page {
315    /// Get the full HTML content of the page including the doctype.
316    ///
317    /// # Example
318    ///
319    /// ```no_run
320    /// # async fn example(page: viewpoint_core::Page) -> Result<(), viewpoint_core::CoreError> {
321    /// let html = page.content().await?;
322    /// println!("Page HTML: {}", html);
323    /// # Ok(())
324    /// # }
325    /// ```
326    ///
327    /// # Errors
328    ///
329    /// Returns an error if the page is closed or evaluation fails.
330    #[instrument(level = "debug", skip(self))]
331    pub async fn content(&self) -> Result<String, PageError> {
332        if self.closed {
333            return Err(PageError::Closed);
334        }
335
336        let result: viewpoint_cdp::protocol::runtime::EvaluateResult = self
337            .connection
338            .send_command(
339                "Runtime.evaluate",
340                Some(EvaluateParams {
341                    expression: r#"
342                        (() => {
343                            const doctype = document.doctype;
344                            const doctypeString = doctype 
345                                ? `<!DOCTYPE ${doctype.name}${doctype.publicId ? ` PUBLIC "${doctype.publicId}"` : ''}${doctype.systemId ? ` "${doctype.systemId}"` : ''}>`
346                                : '';
347                            return doctypeString + document.documentElement.outerHTML;
348                        })()
349                    "#.to_string(),
350                    object_group: None,
351                    include_command_line_api: None,
352                    silent: Some(true),
353                    context_id: None,
354                    return_by_value: Some(true),
355                    await_promise: Some(false),
356                }),
357                Some(&self.session_id),
358            )
359            .await?;
360
361        result
362            .result
363            .value
364            .and_then(|v| v.as_str().map(ToString::to_string))
365            .ok_or_else(|| PageError::EvaluationFailed("Failed to get content".to_string()))
366    }
367
368    /// Set the page HTML content.
369    ///
370    /// Returns a builder for additional options.
371    ///
372    /// # Example
373    ///
374    /// ```no_run
375    /// use viewpoint_core::DocumentLoadState;
376    ///
377    /// # async fn example(page: viewpoint_core::Page) -> Result<(), viewpoint_core::CoreError> {
378    /// // Set simple content
379    /// page.set_content("<html><body>Hello</body></html>").set().await?;
380    ///
381    /// // Set content and wait for network idle
382    /// page.set_content("<html><body><script src='app.js'></script></body></html>")
383    ///     .wait_until(DocumentLoadState::NetworkIdle)
384    ///     .set()
385    ///     .await?;
386    /// # Ok(())
387    /// # }
388    /// ```
389    pub fn set_content(&self, html: impl Into<String>) -> SetContentBuilder<'_> {
390        SetContentBuilder::new(self, html.into())
391    }
392
393    /// Create a builder for injecting script tags.
394    ///
395    /// # Example
396    ///
397    /// ```no_run
398    /// # async fn example(page: viewpoint_core::Page) -> Result<(), viewpoint_core::CoreError> {
399    /// // Add script by URL
400    /// page.add_script_tag().url("https://example.com/script.js").inject().await?;
401    ///
402    /// // Add inline script
403    /// page.add_script_tag().content("console.log('Hello')").inject().await?;
404    ///
405    /// // Add ES6 module
406    /// use viewpoint_core::page::ScriptType;
407    /// page.add_script_tag()
408    ///     .content("export const x = 1;")
409    ///     .script_type(ScriptType::Module)
410    ///     .inject()
411    ///     .await?;
412    /// # Ok(())
413    /// # }
414    /// ```
415    pub fn add_script_tag(&self) -> ScriptTagBuilder<'_> {
416        ScriptTagBuilder::new(self)
417    }
418
419    /// Create a builder for injecting style tags.
420    ///
421    /// # Example
422    ///
423    /// ```no_run
424    /// # async fn example(page: viewpoint_core::Page) -> Result<(), viewpoint_core::CoreError> {
425    /// // Add stylesheet by URL
426    /// page.add_style_tag().url("https://example.com/style.css").inject().await?;
427    ///
428    /// // Add inline CSS
429    /// page.add_style_tag().content("body { background: red; }").inject().await?;
430    /// # Ok(())
431    /// # }
432    /// ```
433    pub fn add_style_tag(&self) -> StyleTagBuilder<'_> {
434        StyleTagBuilder::new(self)
435    }
436}