From 374efd450dfcefb9fbc01ef73b31ea2dd9086b3e Mon Sep 17 00:00:00 2001
From: Eduard Tolosa <edu4rdshl@protonmail.com>
Date: Thu, 9 Feb 2023 18:40:24 -0500
Subject: [PATCH] Forced push due to filesystetm being corrupted. TThis is not
 usable yet.

---
 .rustfmt.toml    |  14 ++++++
 Cargo.toml       |  11 +++++
 src/args.rs      |   1 +
 src/constants.rs |   2 +
 src/lib.rs       | 110 +++++++++++++++++++++++++++++++++++++++++++++++
 5 files changed, 138 insertions(+)
 create mode 100644 .rustfmt.toml
 create mode 100644 Cargo.toml
 create mode 100644 src/args.rs
 create mode 100644 src/constants.rs
 create mode 100644 src/lib.rs

diff --git a/.rustfmt.toml b/.rustfmt.toml
new file mode 100644
index 0000000..fb24a6d
--- /dev/null
+++ b/.rustfmt.toml
@@ -0,0 +1,14 @@
+# Rustfmt configuration file.
+
+# How imports should be grouped into use statements. Imports will be merged or split to the configured level of granularity.
+imports_granularity = "One"
+# Unix or Windows line endings
+newline_style = "Unix"
+# Convert /* */ comments to // comments where possible
+normalize_comments = true
+# Convert #![doc] and #[doc] attributes to //! and /// doc comments.
+normalize_doc_attributes = true
+# Remove nested parens.
+remove_nested_parens = true
+# Reorder impl items. type and const are put first, then macros and methods.
+reorder_impl_items = true
diff --git a/Cargo.toml b/Cargo.toml
new file mode 100644
index 0000000..3b96b09
--- /dev/null
+++ b/Cargo.toml
@@ -0,0 +1,11 @@
+[package]
+name = "rushotter"
+version = "0.1.0"
+edition = "2021"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dependencies]
+futures = "0.3.25"
+thirtyfour = "0.31.0"
+tokio = "1.21.2"
diff --git a/src/args.rs b/src/args.rs
new file mode 100644
index 0000000..8b13789
--- /dev/null
+++ b/src/args.rs
@@ -0,0 +1 @@
+
diff --git a/src/constants.rs b/src/constants.rs
new file mode 100644
index 0000000..74d2e31
--- /dev/null
+++ b/src/constants.rs
@@ -0,0 +1,2 @@
+pub const CHROMEDRIVER_URL: &str = "http://127.0.0.1:9515";
+pub const GECKODRIVER_URL: &str = "http://127.0.0.1:4444";
diff --git a/src/lib.rs b/src/lib.rs
new file mode 100644
index 0000000..058766a
--- /dev/null
+++ b/src/lib.rs
@@ -0,0 +1,110 @@
+use {
+    constants::{CHROMEDRIVER_URL, GECKODRIVER_URL},
+    futures::StreamExt,
+    std::{
+        collections::HashMap,
+        path::{Path, PathBuf},
+    },
+    thirtyfour::prelude::*,
+};
+
+mod constants;
+
+pub async fn take_chrome_screenshots(
+    driver: Option<WebDriver>,
+    chrome_args: Option<&[&str]>,
+    // The first argument is the URL to take a screenshot of and the second argument is the path to save the screenshot to.
+    // If the path is not provided, the screenshot will be saved to the current directory in the format of "url.png".
+    websites: HashMap<String, Option<String>>,
+    threads: Option<usize>,
+) -> WebDriverResult<()> {
+    // Set the number of threads to use. Default to 5.
+    let threads = if let Some(threads) = threads {
+        threads
+    } else {
+        5
+    };
+
+    // Create a new ChromeCapabilities instance.
+    let mut caps = DesiredCapabilities::chrome();
+    caps.set_headless()?;
+
+    // Add any additional Chrome command line arguments.
+    if let Some(chrome_args) = chrome_args {
+        for arg in chrome_args {
+            caps.add_chrome_arg(arg)?;
+        }
+    }
+
+    // Create a new WebDriver instance.
+    let driver = if let Some(driver) = driver {
+        driver
+    } else {
+        WebDriver::new(CHROMEDRIVER_URL, caps).await?
+    };
+
+    println!("Taking screenshots...");
+
+    take_screenshots(websites, &driver, threads).await;
+
+    Ok(())
+}
+
+pub async fn take_firefox_screenshots(
+    driver: Option<WebDriver>,
+    firefox_args: Option<Vec<&str>>,
+    // The first argument is the URL to take a screenshot of and the second argument is the path to save the screenshot to.
+    // If the path is not provided, the screenshot will be saved to the current directory in the format of "url.png".
+    websites: HashMap<String, Option<String>>,
+    threads: usize,
+) -> WebDriverResult<()> {
+    let mut caps = DesiredCapabilities::firefox();
+    caps.set_headless()?;
+
+    if let Some(firefox_args) = firefox_args {
+        for arg in firefox_args {
+            caps.add_firefox_arg(arg)?;
+        }
+    }
+
+    let driver = if let Some(driver) = driver {
+        driver
+    } else {
+        WebDriver::new(GECKODRIVER_URL, caps).await?
+    };
+
+    take_screenshots(websites, &driver, threads).await;
+
+    if driver.quit().await.is_err() {
+        println!("Failed to quit driver");
+    }
+
+    Ok(())
+}
+
+pub async fn take_screenshots(
+    websites: HashMap<String, Option<String>>,
+    driver: &WebDriver,
+    threads: usize,
+) {
+    futures::stream::iter(websites.into_iter().map(|(website, screenshot_path)| {
+        let path = if let Some(screenshot_path) = screenshot_path {
+            Path::new(&screenshot_path).with_extension("png")
+        } else {
+            Path::new(&website).with_extension("png")
+        };
+
+        let fut = screenshot(driver.clone(), website, path);
+
+        tokio::spawn(fut)
+    }))
+    .buffer_unordered(threads)
+    .collect::<Vec<_>>()
+    .await;
+}
+
+pub async fn screenshot(driver: WebDriver, website: String, path: PathBuf) {
+    if driver.goto(&website).await.is_err() || driver.screenshot(&path).await.is_err() {
+        println!("Failed to take screenshot of {}", website);
+    }
+}