From 53987da09f99868f4f8b69a2057ee934a4482010 Mon Sep 17 00:00:00 2001
From: pwnwriter <hey@pwnwriter.xyz>
Date: Wed, 19 Jun 2024 11:46:06 +0545
Subject: [PATCH] feat(args): run arbitary javascript in websites

---
 src/cli/args.rs       |  4 ++++
 src/cli/exec.rs       |  4 ++++
 src/cli/screenshot.rs | 18 ++++++++++++++++++
 3 files changed, 26 insertions(+)

diff --git a/src/cli/args.rs b/src/cli/args.rs
index 08b8ed4..082db2f 100644
--- a/src/cli/args.rs
+++ b/src/cli/args.rs
@@ -76,6 +76,10 @@ pub struct Cli {
     /// Accept invalid certs, trust dns
     #[arg(long)]
     pub accept_invalid_certs: bool,
+
+    /// Run arbiraty javascript
+    #[arg(long)]
+    pub javascript: Option<String>,
 }
 
 #[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, ValueEnum)]
diff --git a/src/cli/exec.rs b/src/cli/exec.rs
index 1fb5b6f..355c2fc 100644
--- a/src/cli/exec.rs
+++ b/src/cli/exec.rs
@@ -29,6 +29,7 @@ pub async fn run(
         screenshot_type,
         ports,
         accept_invalid_certs,
+        javascript,
     }: Cli,
 ) -> anyhow::Result<()> {
     let browser = Path::new(&binary_path);
@@ -105,6 +106,7 @@ pub async fn run(
             fullpage,
             screenshot_type,
             accept_invalid_certs,
+            javascript,
         )
         .await?;
     } else {
@@ -122,6 +124,7 @@ pub async fn run(
                     fullpage,
                     screenshot_type,
                     accept_invalid_certs,
+                    javascript,
                 )
                 .await?;
             }
@@ -142,6 +145,7 @@ pub async fn run(
                     fullpage,
                     screenshot_type,
                     accept_invalid_certs,
+                    javascript,
                 )
                 .await?;
             }
diff --git a/src/cli/screenshot.rs b/src/cli/screenshot.rs
index 382b4a7..5b1b337 100644
--- a/src/cli/screenshot.rs
+++ b/src/cli/screenshot.rs
@@ -23,12 +23,14 @@ pub async fn take_screenshot_in_bulk(
     full_page: bool,
     screenshot_type: ScreenshotType,
     danger_accept_invalid_certs: bool,
+    javascript: Option<String>,
 ) -> anyhow::Result<()> {
     let url_chunks: Vec<Vec<_>> = urls.chunks(tabs).map(ToOwned::to_owned).collect();
     let mut handles = Vec::with_capacity(url_chunks.len());
 
     for urls in url_chunks {
         let browser = Arc::clone(browser);
+        let js = javascript.clone(); // Clone the JavaScript for each chunk
         let handle = tokio::spawn(async move {
             for url in urls {
                 if let Err(error) = take_screenshot(
@@ -40,6 +42,7 @@ pub async fn take_screenshot_in_bulk(
                     full_page,
                     screenshot_type,
                     danger_accept_invalid_certs,
+                    js.clone(), // Pass the JavaScript code
                 )
                 .await
                 {
@@ -67,6 +70,7 @@ pub async fn take_screenshot(
     full_page: bool,
     screenshot_type: ScreenshotType,
     danger_accept_invalid_certs: bool,
+    javascript: Option<String>,
 ) -> anyhow::Result<()> {
     let parsed_url = Url::parse(&url)?;
     let client = reqwest::Client::builder()
@@ -90,6 +94,20 @@ pub async fn take_screenshot(
     };
     let page = browser.new_page(parsed_url.clone()).await?;
     tokio::time::sleep(Duration::from_secs(delay)).await;
+
+
+    // Evaluate JavaScript if provided
+    if let Some(js) = javascript {
+        let result = page.evaluate(js.as_str()).await;
+        match result {
+            Ok(_) => log::info(
+                "JavaScript executed successfully".to_string(),
+                colored::Color::Magenta,
+            ),
+            Err(e) => log::warn(format!("JavaScript execution failed: {:?}", e)),
+        }
+    }
+
     page.save_screenshot(
         ScreenshotParams::builder()
             .format(screenshot_format)