From b3923d7588650af7f45f8f0a2ee29ec20bfbde9e Mon Sep 17 00:00:00 2001 From: Dauaaa Date: Sun, 1 Dec 2024 21:04:27 -0500 Subject: [PATCH 1/2] don't log to stdout if not cli --- src/input.rs | 4 ++- src/lib.rs | 11 +++++++ src/main.rs | 2 ++ src/scanner/mod.rs | 7 ++-- src/scanner/socket_iterator.rs | 2 +- src/tui.rs | 60 +++++++++++++++++++++++++++++----- 6 files changed, 72 insertions(+), 14 deletions(-) diff --git a/src/input.rs b/src/input.rs index a38752118..d8d3e4070 100644 --- a/src/input.rs +++ b/src/input.rs @@ -1,4 +1,6 @@ //! Provides a means to read, parse and hold configuration options for scans. +use crate::print_log; + use clap::{Parser, ValueEnum}; use serde_derive::Deserialize; use std::collections::HashMap; @@ -298,7 +300,7 @@ impl Config { let config: Config = match toml::from_str(&content) { Ok(config) => config, Err(e) => { - println!("Found {e} in configuration file.\nAborting scan.\n"); + print_log!(error, "Found {e} in configuration file.\nAborting scan.\n"); std::process::exit(1); } }; diff --git a/src/lib.rs b/src/lib.rs index e335de58a..7e415fa51 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -56,3 +56,14 @@ pub mod scripts; pub mod address; pub mod generated; + +/// Static variable defining the current state of execution. The cli binary should +/// set it to true by calling set_cli_mode. +#[doc(hidden)] +pub static IS_CLI_MODE: once_cell::sync::OnceCell = once_cell::sync::OnceCell::new(); + +/// Set IS_CLI_MODE to true. +#[doc(hidden)] +pub fn set_cli_mode() { + let _ = IS_CLI_MODE.set(true); +} diff --git a/src/main.rs b/src/main.rs index d79923a13..09fe4cc50 100644 --- a/src/main.rs +++ b/src/main.rs @@ -7,6 +7,7 @@ use rustscan::input::{self, Config, Opts, ScriptsRequired}; use rustscan::port_strategy::PortStrategy; use rustscan::scanner::Scanner; use rustscan::scripts::{init_scripts, Script, ScriptFile}; +use rustscan::set_cli_mode; use rustscan::{detail, funny_opening, output, warning}; use colorful::{Color, Colorful}; @@ -38,6 +39,7 @@ fn main() { #[cfg(not(unix))] let _ = ansi_term::enable_ansi_support(); + set_cli_mode(); env_logger::init(); let mut benchmarks = Benchmark::init(); let mut rustscan_bench = NamedTimer::start("RustScan"); diff --git a/src/scanner/mod.rs b/src/scanner/mod.rs index a4a0fb916..683291eeb 100644 --- a/src/scanner/mod.rs +++ b/src/scanner/mod.rs @@ -1,6 +1,7 @@ //! Core functionality for actual scanning behaviour. use crate::generated::get_parsed_data; use crate::port_strategy::PortStrategy; +use crate::print_log; use log::debug; mod socket_iterator; @@ -290,7 +291,7 @@ impl Scanner { } } Err(e) => { - println!("Err E binding sock {:?}", e); + print_log!(error, "Err E binding sock {:?}", e); Err(e) } } @@ -300,9 +301,9 @@ impl Scanner { fn fmt_ports(&self, socket: SocketAddr) { if !self.greppable { if self.accessible { - println!("Open {socket}"); + print_log!(info, "Open {socket}"); } else { - println!("Open {}", socket.to_string().purple()); + print_log!(info, "Open {}", socket.to_string().purple()); } } } diff --git a/src/scanner/socket_iterator.rs b/src/scanner/socket_iterator.rs index 7813fe71f..a7b602794 100644 --- a/src/scanner/socket_iterator.rs +++ b/src/scanner/socket_iterator.rs @@ -30,7 +30,7 @@ impl<'s> SocketIterator<'s> { } #[allow(clippy::doc_link_with_quotes)] -impl<'s> Iterator for SocketIterator<'s> { +impl Iterator for SocketIterator<'_> { type Item = SocketAddr; /// Returns a socket based on the combination of one of the provided diff --git a/src/tui.rs b/src/tui.rs index 6f215f590..6478044de 100644 --- a/src/tui.rs +++ b/src/tui.rs @@ -5,16 +5,26 @@ #[macro_export] macro_rules! warning { ($name:expr) => { - println!("{} {}", ansi_term::Colour::Red.bold().paint("[!]"), $name); + $crate::print_log!( + warn, + "{} {}", + ansi_term::Colour::Red.bold().paint("[!]"), + $name + ); }; ($name:expr, $greppable:expr, $accessible:expr) => { // if not greppable then print, otherwise no else statement so do not print. if !$greppable { if $accessible { // Don't print the ascii art - println!("{}", $name); + $crate::print_log!(warn, "{}", $name); } else { - println!("{} {}", ansi_term::Colour::Red.bold().paint("[!]"), $name); + $crate::print_log!( + warn, + "{} {}", + ansi_term::Colour::Red.bold().paint("[!]"), + $name + ); } } }; @@ -23,16 +33,26 @@ macro_rules! warning { #[macro_export] macro_rules! detail { ($name:expr) => { - println!("{} {}", ansi_term::Colour::Blue.bold().paint("[~]"), $name); + $crate::print_log!( + info, + "{} {}", + ansi_term::Colour::Blue.bold().paint("[~]"), + $name + ); }; ($name:expr, $greppable:expr, $accessible:expr) => { // if not greppable then print, otherwise no else statement so do not print. if !$greppable { if $accessible { // Don't print the ascii art - println!("{}", $name); + $crate::print_log!(info, "{}", $name); } else { - println!("{} {}", ansi_term::Colour::Blue.bold().paint("[~]"), $name); + $crate::print_log!( + info, + "{} {}", + ansi_term::Colour::Blue.bold().paint("[~]"), + $name + ); } } }; @@ -41,7 +61,8 @@ macro_rules! detail { #[macro_export] macro_rules! output { ($name:expr) => { - println!( + $crate::print_log!( + info, "{} {}", RGansi_term::Colour::RGB(0, 255, 9).bold().paint("[>]"), $name @@ -52,9 +73,10 @@ macro_rules! output { if !$greppable { if $accessible { // Don't print the ascii art - println!("{}", $name); + $crate::print_log!(info, "{}", $name); } else { - println!( + $crate::print_log!( + info, "{} {}", ansi_term::Colour::RGB(0, 255, 9).bold().paint("[>]"), $name @@ -103,3 +125,23 @@ macro_rules! funny_opening { println!("{}\n", random_quote); }; } + +/// Wrapper macro for printing/logging wraps println! and log::$level! +/// 1. if rustscan::IS_CLI_MODE is true calls `println!` +/// 2. if rustscan::IS_CLI_MODE is undefined or false `log::$level!` also sets IS_CLI_MODE +/// to false if it was previously undefined. +/// +/// Library code should call this macro to print information that the binary +/// is expected to print to stdout and library is expected to log at a +/// level specified by parameter $level. +#[doc(hidden)] +#[macro_export] +macro_rules! print_log { + ($level:ident, $($fmt_args:tt)*) => { + if *$crate::IS_CLI_MODE.get_or_init(|| false) { + println!($($fmt_args)*); + } else { + log::$level!($($fmt_args)*); + } + } +} From 42eff43a5b057e2f19b7ec916688bb61bfc8c4a8 Mon Sep 17 00:00:00 2001 From: Dauaaa Date: Tue, 3 Dec 2024 10:39:40 -0500 Subject: [PATCH 2/2] add test and remove once_cell dependency --- Cargo.lock | 3 +-- Cargo.toml | 5 +++- build.rs | 5 ++-- src/generated.rs | 5 ++-- src/lib.rs | 4 +++- tests/log_mode.rs | 28 ++++++++++++++++++++++ tests/log_mode_test_binary/mod.rs | 39 +++++++++++++++++++++++++++++++ 7 files changed, 79 insertions(+), 10 deletions(-) create mode 100644 tests/log_mode.rs create mode 100644 tests/log_mode_test_binary/mod.rs diff --git a/Cargo.lock b/Cargo.lock index 32c761b0e..93914e9b3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1,6 +1,6 @@ # This file is automatically @generated by Cargo. # It is not intended for manual editing. -version = 3 +version = 4 [[package]] name = "addr2line" @@ -1513,7 +1513,6 @@ dependencies = [ "hickory-resolver", "itertools 0.13.0", "log", - "once_cell", "parameterized", "rand", "rlimit", diff --git a/Cargo.toml b/Cargo.toml index 87f27173c..7e4f8c6f1 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -38,7 +38,6 @@ itertools = "0.13.0" hickory-resolver = { version = "0.24.0", features = ["dns-over-rustls"] } anyhow = "1.0.40" text_placeholder = { version = "0.5", features = ["struct_context"] } -once_cell = "1.20.2" [dev-dependencies] parameterized = "2.0.0" @@ -59,6 +58,10 @@ strip = true name = "rustscan" path = "src/main.rs" +[[example]] +name = "log_mode_test_binary" +path = "tests/log_mode_test_binary/mod.rs" + [lints.rust] unexpected_cfgs = { level = "warn", check-cfg = ["cfg(tarpaulin_include)"] } diff --git a/build.rs b/build.rs index 886feb9bb..d29ab448a 100644 --- a/build.rs +++ b/build.rs @@ -62,8 +62,7 @@ fn generate_code(port_payload_map: BTreeMap, Vec>) { let dest_path = PathBuf::from("src/generated.rs"); let mut generated_code = String::new(); - generated_code.push_str("use std::collections::BTreeMap;\n"); - generated_code.push_str("use once_cell::sync::Lazy;\n\n"); + generated_code.push_str("use std::{collections::BTreeMap, sync::LazyLock};\n\n"); generated_code.push_str("fn generated_data() -> BTreeMap, Vec> {\n"); generated_code.push_str(" let mut map = BTreeMap::new();\n"); @@ -92,7 +91,7 @@ fn generate_code(port_payload_map: BTreeMap, Vec>) { generated_code.push_str("}\n\n"); generated_code.push_str( - "static PARSED_DATA: Lazy, Vec>> = Lazy::new(generated_data);\n", + "static PARSED_DATA: LazyLock, Vec>> = LazyLock::new(generated_data);\n", ); generated_code.push_str("pub fn get_parsed_data() -> &'static BTreeMap, Vec> {\n"); generated_code.push_str(" &PARSED_DATA\n"); diff --git a/src/generated.rs b/src/generated.rs index af67ff2d6..11d18aa3f 100644 --- a/src/generated.rs +++ b/src/generated.rs @@ -1,5 +1,4 @@ -use once_cell::sync::Lazy; -use std::collections::BTreeMap; +use std::{collections::BTreeMap, sync::LazyLock}; fn generated_data() -> BTreeMap, Vec> { let mut map = BTreeMap::new(); @@ -2990,7 +2989,7 @@ fn generated_data() -> BTreeMap, Vec> { map } -static PARSED_DATA: Lazy, Vec>> = Lazy::new(generated_data); +static PARSED_DATA: LazyLock, Vec>> = LazyLock::new(generated_data); pub fn get_parsed_data() -> &'static BTreeMap, Vec> { &PARSED_DATA } diff --git a/src/lib.rs b/src/lib.rs index 7e415fa51..9a8c20086 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -41,6 +41,8 @@ //! ``` #![allow(clippy::needless_doctest_main)] +use std::sync::OnceLock; + pub mod tui; pub mod input; @@ -60,7 +62,7 @@ pub mod generated; /// Static variable defining the current state of execution. The cli binary should /// set it to true by calling set_cli_mode. #[doc(hidden)] -pub static IS_CLI_MODE: once_cell::sync::OnceCell = once_cell::sync::OnceCell::new(); +pub static IS_CLI_MODE: OnceLock = OnceLock::new(); /// Set IS_CLI_MODE to true. #[doc(hidden)] diff --git a/tests/log_mode.rs b/tests/log_mode.rs new file mode 100644 index 000000000..3e08e3325 --- /dev/null +++ b/tests/log_mode.rs @@ -0,0 +1,28 @@ +/* + * Test rustscan logging utilities, ensuring library doesn't log to stdout. + */ +use std::{ + io::Read, + process::{Command, Stdio}, +}; + +// need to import here, otherwise cargo thinks I'm not using the test +mod log_mode_test_binary; + +#[test] +fn no_logging_scanner() { + let mut child = Command::new("target/debug/examples/log_mode_test_binary") + .stderr(Stdio::piped()) + .stdout(Stdio::piped()) + .spawn() + .unwrap(); + + child.wait().unwrap(); + + let buf = &mut Vec::new(); + child.stdout.take().unwrap().read_to_end(buf).unwrap(); + assert!(buf.is_empty()); + + child.stderr.take().unwrap().read_to_end(buf).unwrap(); + assert!(buf.is_empty()); +} diff --git a/tests/log_mode_test_binary/mod.rs b/tests/log_mode_test_binary/mod.rs new file mode 100644 index 000000000..ab4650001 --- /dev/null +++ b/tests/log_mode_test_binary/mod.rs @@ -0,0 +1,39 @@ +//! File used just to build a binary for testing if stdout is being logged. +//! Older versions of the library rustscan would write to stdout. This file +//! helps ensure the library only writes to stdout if log is initialized, +//! otherwise it shouldn't write at all. +//! +//! It was necessary to create this file because checking if some code write to +//! stdout is very orthogonal to rust's testing tools. There are utilities but +//! only on unstable rust. This file is used to create a binary that can +//! be executed as child process for testing the behavior. + +#![allow(unused)] + +use std::{net::IpAddr, str::FromStr, time::Duration}; + +use futures::executor::block_on; +use rustscan::{input::ScanOrder, port_strategy::PortStrategy, scanner::Scanner}; + +fn main() { + // "open" tcp connection on random port + let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + // get the port from above connection + let port = listener.local_addr().unwrap().port(); + + // execute + block_on( + Scanner::new( + &[IpAddr::from_str("127.0.0.1").unwrap()], + 100, + Duration::from_secs(5), + 3, + false, + PortStrategy::pick(&None, Some(vec![port]), ScanOrder::Random), + true, + vec![], + false, + ) + .run(), + ); +}