use rswappalyzer_engine::CompiledRuleLibrary;
use rustc_hash::{FxHashMap, FxHashSet};
use std::collections::HashMap;
use std::collections::hash_map::Entry;
use std::hash::BuildHasher;
pub struct DetectionUpdater;
impl DetectionUpdater {
pub fn update<S: BuildHasher>(
detected: &mut HashMap<String, (u8, Option<String>), S>,
tech_name: &str,
confidence: Option<u8>,
version: Option<String>,
) {
let raw_conf = confidence.unwrap_or(100);
let new_version = version;
const NO_VERSION_CONF: u8 = 85;
let new_conf = match &new_version {
Some(_) => raw_conf, None => NO_VERSION_CONF, };
match detected.entry(tech_name.to_string()) {
Entry::Occupied(mut entry) => {
let (old_conf, old_version) = entry.get_mut();
let need_update =
Self::is_new_result_better(new_conf, &new_version, *old_conf, old_version);
if need_update {
*old_conf = new_conf;
*old_version = new_version;
}
}
Entry::Vacant(entry) => {
entry.insert((new_conf, new_version));
}
}
}
pub fn apply_implies<S: BuildHasher>(
compiled_lib: &CompiledRuleLibrary,
detected: &mut HashMap<String, (u8, Option<String>), S>,
) -> (FxHashMap<String, Vec<String>>, Vec<String>) {
let mut imply_source_map: FxHashMap<String, FxHashSet<String>> = FxHashMap::default();
const BASE_IMPLY_CONF: u8 = 90;
const MAX_IMPLY_CONF: u8 = 95;
const BOOST_PER_SOURCE: u8 = 3;
for source_tech_name in detected.keys() {
if let Some(compiled_tech) = compiled_lib.tech_patterns.get(source_tech_name) {
for target_tech_name in &compiled_tech.implies {
let target_tech_name = target_tech_name.trim();
if target_tech_name.is_empty()
|| detected.contains_key(target_tech_name)
{
continue;
}
imply_source_map
.entry(target_tech_name.to_string())
.or_insert_with(FxHashSet::default)
.insert(source_tech_name.to_string());
}
}
}
for (target_tech, source_set) in &imply_source_map {
let source_count = source_set.len() as u8;
let boost = std::cmp::min(source_count * BOOST_PER_SOURCE, MAX_IMPLY_CONF - BASE_IMPLY_CONF);
let final_conf = BASE_IMPLY_CONF + boost;
detected.entry(target_tech.clone()).or_insert((final_conf, None));
}
let mut imply_map = FxHashMap::default();
for (k, v) in imply_source_map {
let mut source_vec = v.into_iter().collect::<Vec<_>>();
source_vec.sort_unstable();
imply_map.insert(k, source_vec);
}
let mut implies_list: Vec<String> = imply_map.keys().cloned().collect();
implies_list.sort_unstable();
(imply_map, implies_list)
}
fn is_new_result_better(
new_conf: u8,
new_version: &Option<String>,
old_conf: u8,
old_version: &Option<String>,
) -> bool {
if new_conf > old_conf {
return true;
}
if new_conf == old_conf {
if old_version.is_none() && new_version.is_some() {
return true;
}
if let (Some(new_ver), Some(old_ver)) = (new_version, old_version) {
return new_ver.len() > old_ver.len();
}
}
false
}
}