From 9ebe429db9c5441801f4c1b132ce2926fd4e6941 Mon Sep 17 00:00:00 2001 From: Alex Jones Date: Tue, 14 Nov 2023 13:38:12 -0500 Subject: [PATCH] feat: large refactor Signed-off-by: Alex Jones --- README.md | 3 +- src/analyze/mod.rs | 148 +++++++++++++++++++----------------- src/analyzer/mod.rs | 9 ++- src/analyzer/sg_analyzer.rs | 18 ++++- 4 files changed, 99 insertions(+), 79 deletions(-) diff --git a/README.md b/README.md index 3db0638..516d02f 100644 --- a/README.md +++ b/README.md @@ -76,4 +76,5 @@ isotope analyze -a S3 - Public instance detection - EBS - Unattached disk - +- SG + - Permissive security group detection \ No newline at end of file diff --git a/src/analyze/mod.rs b/src/analyze/mod.rs index d786ab0..38e6dc2 100644 --- a/src/analyze/mod.rs +++ b/src/analyze/mod.rs @@ -41,39 +41,43 @@ pub async fn run_analysis( let bedrock_client = bedrock::BedrockClient::new(); // Create channels let (tx, rx): (Sender>, Receiver>) = mpsc::channel(); - let analyzers: Vec> = analyzer::generate_analyzers(); - // Progress bars let m = MultiProgress::new(); - + let mut tasks = vec![]; + let mut count = 0; match selected_analyzer { Some(analyzer_arg) => { - let filtered_analyzer = &analyzers - .iter() - .find(|x| x.get_name().as_str() == analyzer_arg); + let filtered_analyzer = analyzer::generate_analyzers() + .into_iter() + .find(|x| x.get_name() == *analyzer_arg); + match filtered_analyzer { Some(x) => { let thread_tx = tx.clone(); - let response = x.run().await; - - match response { - Some(resp_results) => { - thread_tx.send(resp_results).unwrap(); - - } - None => { - thread_tx.send(vec![AnalysisResults::new()]).unwrap(); + let pb = m.add(ProgressBar::new(count)); + tasks.push(tokio::spawn(async move { + pb.inc(1); + let response = x.run().await; + pb.finish_with_message("done..."); + match response { + Some(resp_results) => { + tx.send(resp_results).unwrap(); + pb.finish(); + } + None => { + thread_tx.send(vec![AnalysisResults::new()]).unwrap(); + pb.finish(); + } } - } + })); } None => println!("analyzer of type not found"), } } None => { - let mut tasks = vec![]; + let analyzers: Vec> = analyzer::generate_analyzers(); // Generate threads - let mut count = 0; let alen = analyzers.len(); for current_analyzer in analyzers { let pb = m.add(ProgressBar::new(count)); @@ -99,71 +103,73 @@ pub async fn run_analysis( } None => { thread_tx.send(vec![AnalysisResults::new()]).unwrap(); - + pb.finish(); } } })); - count += 1; - } + count = count + 1; - let mut results: Vec = vec![]; - // Aggregate results - for _n in 0..count { - let rx_result = rx.recv(); - results.append(&mut rx_result.unwrap()); - } - for task in tasks { - task.await.unwrap(); } - m.clear().unwrap(); - let mut processed_results: HashMap> = HashMap::new(); - // generate Vectors aligned to each analyzer type - // Feed results into Bedrock - for mut res in results { - if !res.message.is_empty() { - // Check if the data is in the cache - match conf.fetch_from_cache(&res.message) { - Some(x) => res.advice = x.clone(), - None => { - let result = bedrock_client.enrich(res.message.clone()).await; - // TODO: missing step to copy the bedrock result into res - match result { - Ok(x) => { - res.advice = x.clone(); - // upsert into the cache for next time - conf.clone().upsert_into_cache(&res.message, &x); - // pass ownership over of advice - // check if the processed results analyzer exists as key - // upsert the analysis result into the vector - } - Err(_e) => (), - } - } - } - match processed_results.entry(res.analyzer_name.clone()) { - Entry::Occupied(mut e) => { - e.get_mut().push(res); - } - Entry::Vacant(e) => { - e.insert(vec![res]); + } + } + + let mut results: Vec = vec![]; + // Aggregate results + for _n in 0..tasks.len() { + let rx_result = rx.recv(); + results.append(&mut rx_result.unwrap()); + } + for task in tasks { + task.await.unwrap(); + } + m.clear().unwrap(); + + let mut processed_results: HashMap> = HashMap::new(); + // generate Vectors aligned to each analyzer type + // Feed results into Bedrock + for mut res in results { + if !res.message.is_empty() { + // Check if the data is in the cache + match conf.fetch_from_cache(&res.message) { + Some(x) => res.advice = x.clone(), + None => { + let result = bedrock_client.enrich(res.message.clone()).await; + // TODO: missing step to copy the bedrock result into res + match result { + Ok(x) => { + res.advice = x.clone(); + // upsert into the cache for next time + conf.clone().upsert_into_cache(&res.message, &x); + // pass ownership over of advice + // check if the processed results analyzer exists as key + // upsert the analysis result into the vector } + Err(_e) => (), } } } - - if *enable_json { - let mut p = outputs::Processor::new( - processed_results, - Some(outputs::Configuration::new(*enable_json)), - *explain, - ); - p.print(); - } else { - let mut p = outputs::Processor::new(processed_results, None, *explain); - p.print(); + match processed_results.entry(res.analyzer_name.clone()) { + Entry::Occupied(mut e) => { + e.get_mut().push(res); + } + Entry::Vacant(e) => { + e.insert(vec![res]); + } } } } + + if *enable_json { + let mut p = outputs::Processor::new( + processed_results, + Some(outputs::Configuration::new(*enable_json)), + *explain, + ); + p.print(); + } else { + let mut p = outputs::Processor::new(processed_results, None, *explain); + p.print(); + } Ok(()) } diff --git a/src/analyzer/mod.rs b/src/analyzer/mod.rs index 97e48a1..bb8dcdf 100644 --- a/src/analyzer/mod.rs +++ b/src/analyzer/mod.rs @@ -10,7 +10,7 @@ pub(crate) mod types; mod sg_analyzer; pub fn generate_analyzers() -> Vec> { - vec![ + let analyzers: Vec> = vec![ Box::new(s3_analyzer::S3Analyzer { }), Box::new(sts_analyzer::STSAnalyzer { @@ -19,7 +19,8 @@ pub fn generate_analyzers() -> Vec> { }), Box::new(ebs_analyzer::EbsAnalyzer { }), - // Box::new(sg_analyzer::SecurityGroupsAnalyzer { - // }) - ] + Box::new(sg_analyzer::SecurityGroupsAnalyzer { + }) + ]; + analyzers } diff --git a/src/analyzer/sg_analyzer.rs b/src/analyzer/sg_analyzer.rs index ef54587..4f76904 100644 --- a/src/analyzer/sg_analyzer.rs +++ b/src/analyzer/sg_analyzer.rs @@ -52,10 +52,22 @@ fn has_wide_open_rules(permissions: &Option