From 2924ce4551e5d76fee206e9e3aff9d6737569a03 Mon Sep 17 00:00:00 2001 From: pwnwriter Date: Sat, 2 Sep 2023 20:22:15 +0545 Subject: [PATCH] refactor(entry, args): use seprate dirs // feat(args): stdin support --- Cargo.toml | 3 +- src/ascii.rs | 11 --- src/{ => cli}/args.rs | 18 ++-- src/cli/ascii.rs | 17 ++++ src/cli/mod.rs | 29 +++++++ src/cli/screenshot.rs | 182 ++++++++++++++++++++++++++++++++++++++++ src/log.rs | 27 ++++++ src/main.rs | 188 ++---------------------------------------- 8 files changed, 276 insertions(+), 199 deletions(-) delete mode 100644 src/ascii.rs rename src/{ => cli}/args.rs (68%) create mode 100644 src/cli/ascii.rs create mode 100644 src/cli/mod.rs create mode 100644 src/cli/screenshot.rs create mode 100644 src/log.rs diff --git a/Cargo.toml b/Cargo.toml index 5f84d6a..074a019 100644 --- a/Cargo.toml +++ b/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 diff --git a/src/ascii.rs b/src/ascii.rs deleted file mode 100644 index 62b8e78..0000000 --- a/src/ascii.rs +++ /dev/null @@ -1,11 +0,0 @@ -pub const BAR: &str = r" -──────────────────────────────── -"; - -pub const HXN: &str = r" - ╦ ╦╔═╗╦ ╦╦ ╔═╗╔╗╔ - ╠═╣╠═╣╚╦╝║ 𝖃║ ║║║║ - ╩ ╩╩ ╩ ╩ ╩═╝ ╚═╝╝╚╝v0.1.5 -  Shoot before the blink  - by @PwnWriter - "; diff --git a/src/args.rs b/src/cli/args.rs similarity index 68% rename from src/args.rs rename to src/cli/args.rs index 4ea12e9..7eb1e9c 100644 --- a/src/args.rs +++ b/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, #[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, - #[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, } diff --git a/src/cli/ascii.rs b/src/cli/ascii.rs new file mode 100644 index 0000000..1fbff40 --- /dev/null +++ b/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}") +} diff --git a/src/cli/mod.rs b/src/cli/mod.rs new file mode 100644 index 0000000..2efa6fc --- /dev/null +++ b/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 { + 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 + } +} diff --git a/src/cli/screenshot.rs b/src/cli/screenshot.rs new file mode 100644 index 0000000..a16a4ab --- /dev/null +++ b/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, + outdir: Option, + tabs: Option, + binary_path: String, + width: Option, + height: Option, + timeout: u64, + silent: bool, + stdin: bool, +) -> Result<(), Box> { + 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; // 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, + silent: bool, + timeout_value: u64, +) -> Result<(), Box> { + 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![ + &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(()) +} diff --git a/src/log.rs b/src/log.rs new file mode 100644 index 0000000..84fa646 --- /dev/null +++ b/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); +} diff --git a/src/main.rs b/src/main.rs index 84b4353..5b47934 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,198 +1,24 @@ -mod args; -use args::*; +mod cli; +mod log; +use cli::args; mod colors; -use colors::*; -mod ascii; -use ascii::{BAR, HXN}; -use std::{ - env, - io::{BufRead, BufReader}, - path::Path, - time::Duration, -}; - -use reqwest::get; - -use tokio::{fs, time::timeout}; - -use futures::StreamExt; - -use chromiumoxide::browser::{Browser, BrowserConfig}; -use chromiumoxide::cdp::browser_protocol::page::{ - CaptureScreenshotFormat, CaptureScreenshotParams, -}; -use chromiumoxide::handler::viewport::Viewport; -use chromiumoxide::Page; - use clap::Parser; -use columns::Columns; - #[tokio::main] async fn main() -> Result<(), Box> { - println!("{CYAN}{}{RESET}", HXN); - let cli = Cli::parse(); - run( + let cli = args::Cli::parse(); + crate::cli::screenshot::run( cli.url, Some(cli.outdir), cli.tabs, cli.binary_path, cli.width, cli.height, - cli.timeout_value, + cli.timeout, cli.silent, - ) - .await - .expect("An error occurred while running :("); - - Ok(()) -} - -fn exit_on_error() { - std::process::exit(0); -} - -async fn run( - url: String, - outdir: Option, - tabs: Option, - binary_path: String, - width: Option, - height: Option, - timeout_value: u64, - silent: bool, -) -> Result<(), Box> { - // Check if the browser binary path is valid - if !Path::new(&binary_path).exists() { - println!("{RED}[  ]Browser binary not found at path {}, You should try manually pasing the binary path !!{RESET}", binary_path); - println!("{BLUE}[ ﯦ ]{RESET}{CYAN} $ hxn -b $(which brave) or use --help flag"); - exit_on_error(); - } - 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()?, + cli.stdin, ) .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 = if Path::new(&url).exists() { - // Read URLs from file - let file = std::fs::File::open(&url)?; - let lines = BufReader::new(file).lines().filter_map(Result::ok); - lines.collect() - } else { - // URL is a single URL - vec![url] - }; - - 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); - } - - // Set current working directory to output directory - // So that we can save screenshots in it without specifying whole path. - 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_value)); - handles.push(h); - } - - for handle in handles { - handle - .await? - .expect("Something went wrong while waiting for taking screenshot and saving to file"); - } - - exit_on_error(); - - println!( - "{RED}♥ {GREEN} {YELLOW_BRIGHT}Screenshots saved in dir {outdir}{RED} ♥ {GREEN}{RESET} " - ); - - Ok(()) -} - -async fn take_screenshots( - page: Page, - urls: Vec, - silent: bool, - timeout_value: u64, -) -> Result<(), Box> { - 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![ - &format!("{BLUE}{BAR}"), - &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(()) }