Merge pull request #19 from MrFixThis/refactor/general-app-refactor

refactor: Enhance code quality and cli module

Cargo.toml

@@ -22,11 +22,7 @@ futures = "0.3.28"
 clap = { version = "4.4.4", features = ["derive", "string"] }
 columns = "0.1.0"
 colored = "2.0.4"
-
-[[bin]]
-name = "hxn"
-path = "src/main.rs"
-
+anyhow = "1.0.75"
 
 [profile.dev]
 opt-level = 0

README.md

@@ -105,7 +105,7 @@ Built from the ground up for ease of use, performance, beautiful ui and portabil
      
 
   ```bash
-  hxn -b $(which brave) -u urls.txt
+  hxn -b $(which brave) -f urls.txt
   ```
 
   ![many](https://github.com/pwnwriter/haylxon/assets/90331517/86f987d9-0961-4247-841e-18aee6aaf53f)

src/cli/args.rs

@@ -1,44 +1,101 @@
-use crate::cli::splash;
-use clap::Parser;
+use super::ascii;
+use clap::{Args, Parser};
 
-#[derive(Parser)]
-#[command(author, version, about = splash()  )]
-#[command(propagate_version = true)]
-#[command(arg_required_else_help = true)]
-pub struct Cli {
-    #[arg(required = false, short, long)]
-    /// a single url or a file containing multiple urls
+#[derive(Debug, Args)]
+#[group(required = false, multiple = false, conflicts_with = "stdin")]
+pub struct Input {
+    /// Website URL
+    #[arg(short, long)]
     pub url: Option<String>,
 
-    #[arg(short, long, default_value = "hxnshots")]
+    /// Path of the file containing URLs
+    #[arg(short, long)]
+    pub file_path: Option<String>,
+}
+
+#[derive(Debug, Parser)]
+#[command(
+    author,
+    version,
+    about = ascii::splash(),
+    propagate_version = true,
+    arg_required_else_help = true
+)]
+pub struct Cli {
+    /// Browser binary path
+    #[arg(short, long, default_value = "/usr/bin/chrome")]
+    pub binary_path: String,
+
+    /// Read urls from the standard input
+    #[arg(long)]
+    pub stdin: bool,
+
+    #[command(flatten)]
+    pub input: Input,
+
     /// Output directory to save screenshots
+    #[arg(short, long, default_value = "hxnshots")]
     pub outdir: String,
 
-    #[arg(short, long, default_value = "4")]
     /// Maximum number of parallel tabs
-    pub tabs: Option<usize>,
-
-    #[arg(short, long, default_value = "/usr/bin/google-chrome")]
-    /// Browser binary path
-    pub binary_path: String,
+    #[arg(short, long, default_value = "4")]
+    pub tabs: usize,
 
-    #[arg(short = 'x', long, default_value = "1440")]
     /// Width of the website // URL
-    pub width: Option<u32>,
+    #[arg(short = 'x', long, default_value = "1440")]
+    pub width: u32,
 
-    #[arg(short = 'y', long, default_value = "900")]
     /// Height of the website // URL
-    pub height: Option<u32>,
+    #[arg(short = 'y', long, default_value = "900")]
+    pub height: u32,
 
-    #[arg(long, default_value = "10")]
     /// Define timeout for urls
+    #[arg(long, default_value = "10")]
     pub timeout: u64,
 
-    #[arg(long)]
     /// Silent mode (suppress all console output)
+    #[arg(long)]
     pub silent: bool,
+}
 
-    #[arg(long)]
-    /// Read urls from the standard in
-    pub stdin: bool,
+#[cfg(test)]
+mod tests {
+    use clap::error::ErrorKind;
+
+    use super::*;
+
+    #[test]
+    fn test_no_input_urls() {
+        let args = Cli::try_parse_from(["-b my_browser"]);
+        assert!(args.is_err());
+        assert_eq!(
+            args.unwrap_err().kind(),
+            ErrorKind::DisplayHelpOnMissingArgumentOrSubcommand
+        );
+    }
+
+    #[test]
+    fn test_stdin_source_mutual_exclusion() {
+        let args = Cli::try_parse_from([
+            "-b my_browser",
+            "--stdin",
+            "-u https://example.com",
+            "-f /my/file",
+        ]);
+        assert!(args.is_err());
+        assert_eq!(args.unwrap_err().kind(), ErrorKind::ArgumentConflict);
+    }
+
+    #[test]
+    fn test_url_mutual_exclusion_with_file_path() {
+        let args = Cli::try_parse_from(["-b my_browser", "-u https://example.com", "-f /my/file"]);
+        assert!(args.is_err());
+        assert_eq!(args.unwrap_err().kind(), ErrorKind::ArgumentConflict);
+    }
+
+    #[test]
+    fn test_file_path_as_source() {
+        let args = Cli::try_parse_from(["-b my_browser", "-f /my/file"]);
+        assert!(args.is_ok());
+    }
 }

src/cli/mod.rs

@@ -1,52 +1,28 @@
 pub mod args;
-pub mod ascii;
+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
+    use std::{io, path::Path};
+
+    /// Reads user's input from stdin line by line.
+    #[inline]
+    pub fn read_urls_from_stdin() -> anyhow::Result<Vec<String>> {
+        Ok(io::read_to_string(io::stdin().lock())?
+            .lines()
+            .map(|url| url.trim().to_owned())
+            .collect())
     }
 
-    #[allow(dead_code)]
-    pub fn read_urls_from_file(url: &Option<String>) -> Vec<String> {
-        let mut urls = Vec::new();
-
-        if let Some(url) = url {
-            if std::path::Path::new(url).exists() {
-                if let Ok(file) = std::fs::File::open(url) {
-                    let lines = std::io::BufReader::new(file).lines().map_while(Result::ok);
-                    urls = lines.collect();
-                } else {
-                    urls = vec![url.clone()];
-                }
-            } else {
-                urls = vec![url.clone()];
-            }
-        }
-
-        urls
+    /// Reads URLs from a file.
+    #[inline]
+    pub fn read_urls_from_file<T: AsRef<Path>>(file_path: T) -> anyhow::Result<Vec<String>> {
+        Ok(std::fs::read_to_string(file_path)?
+            .lines()
+            .map(|url| url.trim().to_owned())
+            .collect())
     }
 
-    #[allow(dead_code)]
+    #[allow(unused)]
     pub fn read_ports() {}
 }