Rust/gscan

rust/examples/gscan/Cargo.toml

@@ -1,16 +1,17 @@
 [package]
 name = "gscan"
-version = "1.0.0"
+version = "1.1.0"
 edition = "2021"
 authors = ["GDATA CyberDefense AG <[email protected]>"]
 license = "MIT"
 description = "GDATA Verdict-as-a-Service CLI Scanner"
 publish = false
 
 [dependencies]
-vaas = { version = "6.0.0" }
+vaas = { version = "6.1.1" }
 tokio = { version = "1.37", features = ["rt-multi-thread", "macros"] }
-clap = { version = "4.5.4", features = ["env", "cargo"] }
+clap = { version = "4.5.18", features = ["env", "cargo", "derive"] }
 reqwest = "0.12.4"
 futures = "0.3.30"
 dotenv = "0.15"
+walkdir = "2.5.0"

rust/examples/gscan/src/main.rs

@@ -1,118 +1,132 @@
-use clap::{crate_authors, crate_description, crate_name, crate_version, Arg, ArgAction, Command};
-use reqwest::Url;
-use std::{collections::HashMap, path::PathBuf, str::FromStr};
-use vaas::{
-    auth::authenticators::ClientCredentials, error::VResult, CancellationToken, Connection, Vaas,
-    VaasVerdict,
-};
-
-#[tokio::main]
-async fn main() -> VResult<()> {
-    let matches = Command::new(crate_name!())
-        .version(crate_version!())
-        .author(crate_authors!())
-        .about(crate_description!())
-        .arg(
-            Arg::new("files")
-                .short('f')
-                .long("files")
-                .required_unless_present("urls")
-                .action(ArgAction::Append)
-                .help("List of files to scan separated by whitepace"),
-        )
-        .arg(
-            Arg::new("urls")
-                .short('u')
-                .long("urls")
-                .action(ArgAction::Append)
-                .required_unless_present("files")
-                .help("List of urls to scan separated by whitepace"),
-        )
-        .arg(
-            Arg::new("client_id")
-                .short('i')
-                .long("client_id")
-                .env("CLIENT_ID")
-                .action(ArgAction::Set)
-                .help("Set your vaas username"),
-        )
-        .arg(
-            Arg::new("client_secret")
-                .short('s')
-                .long("client_secret")
-                .env("CLIENT_SECRET")
-                .action(ArgAction::Set)
-                .help("Set your vaas password"),
-        )
-        .get_matches();
-
-    let files = matches
-        .get_many::<String>("files")
-        .unwrap_or_default()
-        .map(|f| PathBuf::from_str(f).unwrap_or_else(|_| panic!("Not a valid file path: {}", f)))
-        .collect::<Vec<PathBuf>>();
-
-    let urls = matches
-        .get_many::<String>("urls")
-        .unwrap_or_default()
-        .map(|f| Url::parse(f).unwrap_or_else(|_| panic!("Not a valid url: {}", f)))
-        .collect::<Vec<Url>>();
-
-    let client_id = matches
-        .get_one::<String>("client_id")
-        .expect("--client_id or the enviroment variable CLIENT_ID must be set");
-    let client_secret = matches
-        .get_one::<String>("client_secret")
-        .expect("--client_secret or the enviroment variable CLIENT_SECRET must be set");
-
-    let authenticator = ClientCredentials::new(client_id.to_owned(), client_secret.to_owned());
-    let vaas_connection = Vaas::builder(authenticator).build()?.connect().await?;
-
-    let file_verdicts = scan_files(&files, &vaas_connection).await?;
-    let url_verdicts = scan_urls(&urls, &vaas_connection).await?;
-
-    file_verdicts
-        .iter()
-        .for_each(|(f, v)| print_verdicts(f.display().to_string(), v));
-
-    url_verdicts.iter().for_each(|(u, v)| print_verdicts(u, v));
-
-    Ok(())
-}
-
-fn print_verdicts<I: AsRef<str>>(i: I, v: &VResult<VaasVerdict>) {
-    print!("{} -> ", i.as_ref());
-    match v {
-        Ok(v) => {
-            println!("{}", v.verdict);
-        }
-        Err(e) => {
-            println!("{}", e.to_string());
-        }
-    };
-}
-
-async fn scan_files<'a>(
-    files: &'a [PathBuf],
-    vaas_connection: &Connection,
-) -> VResult<Vec<(&'a PathBuf, VResult<VaasVerdict>)>> {
-    let ct = CancellationToken::from_minutes(1);
-    let verdicts = vaas_connection.for_file_list(files, &ct).await;
-    let results = files.iter().zip(verdicts).collect();
-
-    Ok(results)
-}
-
-async fn scan_urls(
-    urls: &[Url],
-    vaas_connection: &Connection,
-) -> VResult<HashMap<Url, Result<VaasVerdict, vaas::error::Error>>> {
-    let ct = CancellationToken::from_minutes(1);
-    let mut verdicts = HashMap::new();
-    for url in urls {
-        let verdict = vaas_connection.for_url(url, &ct).await;
-        verdicts.insert(url.to_owned(), verdict);
-    }
-
-    Ok(verdicts)
-}
+use clap::Parser;
+use dotenv::dotenv;
+use futures::{stream, StreamExt};
+use reqwest::Url;
+use std::{collections::HashMap, path::PathBuf};
+use vaas::{
+    auth::authenticators::ClientCredentials, error::VResult, CancellationToken, Connection, Vaas,
+    VaasVerdict,
+};
+use walkdir::WalkDir;
+
+#[derive(Parser, Debug)]
+#[command(version, about, long_about = None)]
+struct Args {
+    #[arg(
+        short = 'i',
+        long = "client_id",
+        env = "CLIENT_ID",
+        help = "Set your VaaS client ID"
+    )]
+    client_id: String,
+
+    #[arg(
+        short = 's',
+        long = "client_secret",
+        env = "CLIENT_SECRET",
+        help("Set your VaaS client secret")
+    )]
+    client_secret: String,
+
+    #[arg(long, help = "Lookup the SHA256 hash")]
+    use_hash_lookup: bool,
+
+    #[arg(long, help = "Use the cache")]
+    use_cache: bool,
+
+    #[arg(
+        short = 'f',
+        long,
+        num_args=1..,
+        required_unless_present("urls"),
+        help = "List of files to scan separated by whitepace"
+    )]
+    files: Vec<PathBuf>,
+
+    #[arg(
+        short = 'u',
+        long,
+        num_args=1..,
+        required_unless_present("files"),
+        help = "List of urls to scan separated by whitepace"
+    )]
+    urls: Vec<Url>,
+}
+
+
+#[tokio::main]
+async fn main() -> VResult<()> {
+    dotenv().ok();
+    let args = Args::parse();
+
+    let files = expand_directories(&args.files);
+
+    let authenticator = ClientCredentials::new(args.client_id.clone(), args.client_secret.clone());
+    let vaas_connection = Vaas::builder(authenticator)
+        .use_hash_lookup(args.use_hash_lookup)
+        .use_cache(args.use_cache)
+        .build()?
+        .connect()
+        .await?;
+
+    let files = stream::iter(files);
+    let vaas_reference = &vaas_connection;
+    files.for_each_concurrent(2,|p| async move {
+            let ct = CancellationToken::from_minutes(1);
+            let verdict = vaas_reference.for_file(&p, &ct).await;
+            print_verdicts(p.display().to_string(), &verdict);
+    }).await;
+
+    let url_verdicts = scan_urls(args.urls.as_ref(), &vaas_connection).await?;
+
+    url_verdicts.iter().for_each(|(u, v)| {
+        print_verdicts(u, v);
+    });
+
+    Ok(())
+}
+
+fn print_verdicts<I: AsRef<str>>(i: I, v: &VResult<VaasVerdict>) {
+    print!("{} -> ", i.as_ref());
+    match v {
+        Ok(v) => {
+            println!("{}", v.verdict);
+        }
+        Err(e) => {
+            println!("{}", e);
+        }
+    };
+}
+
+
+async fn scan_urls(
+    urls: &[Url],
+    vaas_connection: &Connection,
+) -> VResult<HashMap<Url, Result<VaasVerdict, vaas::error::Error>>> {
+    let ct = CancellationToken::from_minutes(1);
+    let mut verdicts = HashMap::new();
+    for url in urls {
+        let verdict = vaas_connection.for_url(url, &ct).await;
+        verdicts.insert(url.to_owned(), verdict);
+    }
+
+    Ok(verdicts)
+}
+
+fn expand_directories(files: &[PathBuf]) -> impl Iterator<Item = PathBuf> + '_ {
+    files.iter().flat_map(expand_entry)
+}
+
+fn expand_entry(p: &PathBuf) -> Box<dyn Iterator<Item = PathBuf>> {
+    if p.is_file() {
+        return Box::new(std::iter::once(p.clone()));
+    }
+
+    let files_in_directory = WalkDir::new(p)
+        .into_iter()
+        .filter_map(|e| e.ok())
+        .filter(|e| e.file_type().is_file())
+        .map(|e| e.path().to_path_buf());
+
+    Box::new(files_in_directory)
+}