refactor(entry, args): use seprate dirs // feat(args): stdin support

Cargo.toml

@@ -19,8 +19,9 @@ reqwest = "0.11.16"
 tokio = { version = "1.27.0", features = ["full"] }
 chromiumoxide = {  version = "0.5.0", features = ["tokio-runtime"], default-features = false }
 futures = "0.3.28"
-clap = { version = "4.2.4", features = ["derive"] }
+clap = { version = "4.2.4", features = ["derive", "string"] }
 columns = "0.1.0"
+colored = "2.0.4"
 
 [profile.dev]
 opt-level = 0

src/args.rs → src/cli/args.rs

@@ -1,11 +1,14 @@
+use crate::cli::splash;
 use clap::Parser;
 
 #[derive(Parser)]
-#[command(author, version, about, long_about = "Grab screenshots of your domain list right from terminal.")]
+#[command(author, version, about = splash()  )]
+#[command(propagate_version = true)]
+#[command(arg_required_else_help = true)]
 pub struct Cli {
-    #[arg(short, long)]
+    #[arg(required = false,short, long)]
     /// Website URL/filename of file containing URLs
-    pub url: String,
+    pub url: Option<String>,
 
     #[arg(short, long, default_value = "hxnshots")]
     /// Output directory to save screenshots
@@ -27,12 +30,15 @@ pub struct Cli {
     /// Height of the website // URL
     pub height: Option<u32>,
 
-    #[arg(short = 'k', long, default_value = "10")]
+    #[arg(long, default_value = "10")]
     /// Define timeout for urls
-    pub timeout_value: u64,
+    pub timeout: u64,
 
-    #[arg(short, long)]
+    #[arg(long)]
     /// Silent mode (suppress all console output)
     pub silent: bool,
 
+    #[arg(long)]
+    /// Read urls from the standard in
+    pub stdin: bool,
 }

src/cli/ascii.rs

@@ -0,0 +1,17 @@
+use colored::Colorize;
+// pub const BAR: &str = r"
+// ────────────────────────────────
+// ";
+
+pub fn splash() -> String {
+    let logo = r"
+ ╦ ╦╔═╗╦ ╦╦   ╔═╗╔╗╔
+ ╠═╣╠═╣╚╦╝║  𝖃║ ║║║║
+ ╩ ╩╩ ╩ ╩ ╩═╝ ╚═╝╝╚╝v0.1.5
+             by @PwnWriter
+ "
+    .purple();
+    let quote = " Shoot before the blink  ".italic();
+
+    format!("{logo} {quote}")
+}

src/cli/mod.rs

@@ -0,0 +1,29 @@
+pub mod args;
+pub mod ascii;
+pub mod screenshot;
+pub use args::*;
+pub use ascii::*;
+pub use screenshot::*;
+
+pub mod hxn_helper {
+
+    use std::io::BufRead;
+
+    /// https://www.youtube.com/watch?v=K_wnB9ibCMg&t=1078s
+    /// Reads user input from stdin line by line
+    pub fn read_urls_from_stdin() -> Vec<String> {
+        let mut input = String::new();
+        let mut urls = Vec::new();
+
+        loop {
+            input.clear();
+            match std::io::stdin().lock().read_line(&mut input) {
+                Ok(0) => break, // EOF reached
+                Ok(_) => urls.push(input.trim().to_string()),
+                Err(err) => panic!("Error reading from stdin: {}", err),
+            }
+        }
+
+        urls
+    }
+}

src/cli/screenshot.rs

@@ -0,0 +1,182 @@
+use crate::colors::*;
+use crate::log::error;
+use chromiumoxide::browser::{Browser, BrowserConfig};
+use chromiumoxide::handler::viewport::Viewport;
+use futures::StreamExt;
+use std::{
+    env,
+    io::{BufRead, BufReader},
+    path::Path,
+};
+use tokio::{fs, time::timeout};
+
+use chromiumoxide::cdp::browser_protocol::page::{
+    CaptureScreenshotFormat, CaptureScreenshotParams,
+};
+use chromiumoxide::Page;
+use columns::Columns;
+use core::time::Duration;
+use reqwest::get;
+
+#[allow(clippy::too_many_arguments)]
+pub async fn run(
+    url: Option<String>,
+    outdir: Option<String>,
+    tabs: Option<usize>,
+    binary_path: String,
+    width: Option<u32>,
+    height: Option<u32>,
+    timeout: u64,
+    silent: bool,
+    stdin: bool,
+) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
+    if !Path::new(&binary_path).exists() {
+        error("Unble to locate browser binary");
+
+        std::process::exit(0);
+    }
+    let outdir = match outdir {
+        Some(dir) => dir,
+        None => "hxnshots".to_string(),
+    };
+
+    let viewport_width = width.unwrap_or(1440);
+    let viewport_height = height.unwrap_or(900);
+
+    let (browser, mut handler) = Browser::launch(
+        BrowserConfig::builder()
+            .no_sandbox()
+            .window_size(viewport_width, viewport_height)
+            .chrome_executable(Path::new(&binary_path))
+            .viewport(Viewport {
+                width: viewport_width,
+                height: viewport_height,
+                device_scale_factor: None,
+                emulating_mobile: false,
+                is_landscape: false,
+                has_touch: false,
+            })
+            .build()?,
+    )
+    .await?;
+
+    let _handle = tokio::task::spawn(async move {
+        loop {
+            let _ = handler.next().await;
+        }
+    });
+
+    if fs::metadata(&outdir).await.is_err() {
+        fs::create_dir(&outdir).await?;
+    }
+
+    let urls: Vec<String>; // Define the 'urls' variable outside the match statement
+
+    #[allow(unreachable_patterns)]
+    match stdin {
+        true => {
+            urls = crate::cli::hxn_helper::read_urls_from_stdin();
+        }
+
+        false => {
+            if let Some(url) = &url {
+                if Path::new(url).exists() {
+                    // Read URLs from file
+                    let file = std::fs::File::open(url)?;
+                    let lines = BufReader::new(file).lines().map_while(Result::ok);
+                    urls = lines.collect(); // Assign the collected lines to 'urls'
+                } else {
+                    // URL is a single URL
+                    urls = vec![url.clone()]; // Assign the single URL to 'urls'
+                }
+            } else {
+                // Handle the case where 'url' is None (you can decide what to do in this case)
+                // For now, let's assume it's an empty vector
+                urls = vec![];
+            }
+        }
+
+        _ => {
+            // Handle other cases if needed
+            // For now, let's assume it's an empty vector
+            urls = vec![];
+        }
+    }
+
+    let mut url_chunks = Vec::new();
+
+    for chunk in urls.chunks(tabs.unwrap_or(4)) {
+        let mut urls = Vec::new();
+        for url in chunk {
+            if let Ok(url) = url::Url::parse(url) {
+                urls.push(url);
+            }
+        }
+        url_chunks.push(urls);
+    }
+
+    env::set_current_dir(Path::new(&outdir))?;
+
+    let mut handles = Vec::new();
+
+    for chunk in url_chunks {
+        let n_tab = browser.new_page("about:blank").await?;
+        let h = tokio::spawn(take_screenshots(n_tab, chunk, silent, timeout));
+        handles.push(h);
+    }
+
+    for handle in handles {
+        handle
+            .await?
+            .expect("Something went wrong while waiting for taking screenshot and saving to file");
+    }
+
+    println!("Screenshots saved in dir {outdir}");
+
+    Ok(())
+}
+
+async fn take_screenshots(
+    page: Page,
+    urls: Vec<reqwest::Url>,
+    silent: bool,
+    timeout_value: u64,
+) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
+    for url in urls {
+        let url = url.as_str();
+        if let Ok(Ok(_res)) = timeout(Duration::from_secs(timeout_value), get(url)).await {
+            let filename = url.replace("://", "-").replace('/', "_") + ".png";
+            page.goto(url)
+                .await?
+                .save_screenshot(
+                    CaptureScreenshotParams::builder()
+                        .format(CaptureScreenshotFormat::Png)
+                        .build(),
+                    filename,
+                )
+                .await?;
+
+            let _info = Columns::from(vec![
+                format!("{RESET}").split('\n').collect::<Vec<&str>>(),
+                vec![
+                    &format!("{BLUE}"),
+                    &format!("{GREEN}[{CYAN}  {GREEN}] URL={GREEN}{}", url),
+                    &format!(
+                        "{BLUE}[{CYAN}  {YELLOW}] Title={GREEN}{}",
+                        page.get_title().await?.unwrap_or_default()
+                    ),
+                    &format!("{BLUE}[{CYAN} ﯜ {YELLOW}] Status={GREEN}{}", _res.status()),
+                ],
+            ])
+            .set_tabsize(0)
+            .make_columns();
+            if !silent {
+                println!("{_info}");
+            }
+        } else {
+            println!("{RED}[-] Timed out URL = {YELLOW}{}", url);
+        }
+    }
+
+    Ok(())
+}

src/log.rs

@@ -0,0 +1,27 @@
+use colored::{Color, Colorize};
+
+/// Prints the given message to the console and aborts the process.
+#[allow(dead_code)]
+pub fn abort(msg: &str) -> ! {
+    error(msg);
+    std::process::exit(1);
+}
+
+#[allow(dead_code)]
+pub fn info(msg: &str, color: Color) {
+    println!("{}: {}", "info".bold().color(color), msg);
+}
+
+pub fn error(msg: &str) {
+    println!("{}: {}", "error".bold().color(Color::Red), msg);
+}
+
+#[allow(dead_code)]
+pub fn success(msg: &str) {
+    println!("{}: {}", "success".bold().color(Color::Green), msg);
+}
+
+#[allow(dead_code)]
+pub fn warn(msg: &str) {
+    println!("{}: {}", "warning".bold().color(Color::Yellow), msg);
+}