use crate::extractor::{extract_manual_features, ngrams_for_bucket};
use crate::model::Model;
use crate::predictor::{predict_tree_with_path, predict_url};
use std::collections::HashMap;
const MAX_INDICATORS: usize = 5;
const MIN_TREE_VOTES: usize = 3;
#[derive(Debug, Clone)]
pub struct Prediction {
pub score: f32,
pub indicators: Vec<Indicator>,
}
#[derive(Debug, Clone)]
pub struct Indicator {
pub category: IndicatorCategory,
pub description: String,
pub weight: f32,
pub source: IndicatorSource,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum IndicatorCategory {
Domain,
Path,
Structure,
NGram,
}
#[derive(Debug, Clone)]
pub enum IndicatorSource {
Model {
tree_count: usize,
total_trees: usize,
},
Heuristic,
}
pub fn predict_url_detailed(url: &str, model: &Model) -> Prediction {
let score = predict_url(url, model);
let total_trees = model.trees.len();
let features = model.extract_features(url);
let manual_features = extract_manual_features(url);
let mut feature_votes: HashMap<i32, usize> = HashMap::new();
for tree in &model.trees {
let (_, path) = predict_tree_with_path(tree, &features);
for step in &path {
*feature_votes.entry(step.feature_idx).or_insert(0) += 1;
}
}
let mut ranked_features: Vec<(i32, usize)> = feature_votes
.into_iter()
.filter(|&(_, count)| count >= MIN_TREE_VOTES)
.collect();
ranked_features.sort_by_key(|b| std::cmp::Reverse(b.1));
let mut indicators: Vec<Indicator> = Vec::new();
let mut covered_manual: Vec<usize> = Vec::new();
for (feature_idx, tree_count) in &ranked_features {
if indicators.len() >= MAX_INDICATORS {
break;
}
let feat_idx = *feature_idx;
if (feat_idx as usize) >= model.n_features {
let manual_idx = feat_idx as usize - model.n_features;
if manual_idx >= manual_features.len() {
continue;
}
let val = manual_features[manual_idx];
if let Some((category, desc)) = manual_feature_indicator(manual_idx, val, url) {
covered_manual.push(manual_idx);
indicators.push(Indicator {
category,
description: desc,
weight: *tree_count as f32 / total_trees.max(1) as f32,
source: IndicatorSource::Model {
tree_count: *tree_count,
total_trees,
},
});
}
}
}
for (idx, &val) in manual_features.iter().enumerate() {
if indicators.len() >= MAX_INDICATORS {
break;
}
if covered_manual.contains(&idx) {
continue;
}
if let Some((category, desc)) = manual_feature_indicator(idx, val, url) {
covered_manual.push(idx);
indicators.push(Indicator {
category,
description: desc,
weight: 0.5,
source: IndicatorSource::Heuristic,
});
}
}
for (feature_idx, tree_count) in &ranked_features {
if indicators.len() >= MAX_INDICATORS {
break;
}
let feat_idx = *feature_idx;
if (feat_idx as usize) < model.n_features {
let bucket = feat_idx as usize;
let ngrams = ngrams_for_bucket(url, bucket, model.n_features, model.ngram_range);
if let Some(best) = ngrams.iter().max_by_key(|g| g.len()) {
indicators.push(Indicator {
category: IndicatorCategory::NGram,
description: format!("Suspicious character pattern: '{}'", best),
weight: *tree_count as f32 / total_trees.max(1) as f32,
source: IndicatorSource::Model {
tree_count: *tree_count,
total_trees,
},
});
}
}
}
Prediction { score, indicators }
}
fn manual_feature_indicator(
idx: usize,
val: f32,
url: &str,
) -> Option<(IndicatorCategory, String)> {
match idx {
2 if val >= 4.0 => Some((
IndicatorCategory::Structure,
"URL length abnormal".to_string(),
)),
3 if val >= 3.0 => Some((
IndicatorCategory::Domain,
"Domain length abnormal".to_string(),
)),
4 if val == 1.0 => Some((
IndicatorCategory::Domain,
"Uses IP address instead of domain name".to_string(),
)),
5 if val >= 3.0 => Some((
IndicatorCategory::Domain,
"Excessive subdomain levels".to_string(),
)),
7 if val >= 3.0 => Some((
IndicatorCategory::Structure,
"Excessive hyphens in URL".to_string(),
)),
9 if val >= 1.0 => Some((
IndicatorCategory::Structure,
"Contains @ symbol (URL obfuscation)".to_string(),
)),
10 if val >= 1.0 => Some((
IndicatorCategory::Structure,
"High percent-encoding usage".to_string(),
)),
14 if val > 0.3 => Some((
IndicatorCategory::Structure,
"High digit-to-character ratio".to_string(),
)),
15 if val >= 3.0 => Some((IndicatorCategory::Path, "Path length abnormal".to_string())),
18 if val == 1.0 => {
let keyword = find_sensitive_keyword(url);
Some((
IndicatorCategory::Path,
if let Some(kw) = keyword {
format!("Contains sensitive keyword: '{}'", kw)
} else {
"Contains sensitive keyword (login/verify/paypal)".to_string()
},
))
}
19 if val >= 0.5 => Some((
IndicatorCategory::Domain,
format!(
"Domain closely resembles a known brand (impersonation score {:.2})",
val
),
)),
20 if val >= 0.7 => Some((
IndicatorCategory::Domain,
format!("Random / high-entropy subdomain (entropy {:.2})", val),
)),
21 if val >= 5.0 => Some((
IndicatorCategory::Structure,
format!("Deep path structure ({} levels)", val as i32),
)),
22 if val >= 5.0 => Some((
IndicatorCategory::Structure,
format!("Excessive query parameters ({})", val as i32),
)),
23 if val == 1.0 => Some((
IndicatorCategory::Domain,
"Explicit port number in URL".to_string(),
)),
24 if val >= 1.0 => Some((
IndicatorCategory::Domain,
"Non-standard port number (>= 1000)".to_string(),
)),
25 if val > 0.5 => Some((
IndicatorCategory::Structure,
format!("High hexadecimal character ratio ({:.2})", val),
)),
26 if val < 0.3 => Some((
IndicatorCategory::Structure,
format!("Low alphabetic character ratio ({:.2})", val),
)),
27 if val > 0.1 => Some((
IndicatorCategory::Structure,
format!("High uppercase ratio (obfuscation) ({:.2})", val),
)),
28 if val >= 20.0 => Some((
IndicatorCategory::Domain,
format!("Abnormally long domain label ({} chars)", val as i32),
)),
29 if val >= 5.0 => Some((
IndicatorCategory::Domain,
format!("Excessive domain label count ({})", val as i32),
)),
30 if val >= 15.0 => Some((
IndicatorCategory::Domain,
format!("Long average domain label length ({:.1})", val),
)),
31 if val > 0.3 => Some((
IndicatorCategory::Domain,
format!("High digit ratio in domain ({:.2})", val),
)),
32 if val >= 0.7 => Some((
IndicatorCategory::Domain,
format!("High domain entropy ({:.2})", val),
)),
35 if val == 1.0 => Some((
IndicatorCategory::Structure,
"Double-slash obfuscation in path".to_string(),
)),
36 if val == 1.0 => Some((
IndicatorCategory::Domain,
"Domain label ends with hyphen".to_string(),
)),
37 if val == 1.0 => Some((
IndicatorCategory::Domain,
"Domain label starts with digit".to_string(),
)),
38 if val == 1.0 => Some((
IndicatorCategory::Domain,
"Dangerous URL scheme (data:/javascript:)".to_string(),
)),
_ => None,
}
}
fn find_sensitive_keyword(url: &str) -> Option<String> {
let normalized = crate::extractor::normalize_url(url);
let url_lower = normalized.to_lowercase();
let re = crate::extractor::sensitive_word_regex();
re.find(&url_lower).map(|m| m.as_str().to_string())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_predict_url_detailed_phishing() {
let model = crate::load_default_model().expect("Failed to load model");
let result = predict_url_detailed(
"nobell.it/70ffb52d079109dca5664cce6f317373782/login.SkyPe.com",
&model,
);
assert!(
result.score >= 0.20,
"Phishing URL should score >= 0.20, got {}",
result.score
);
assert!(
!result.indicators.is_empty(),
"Phishing URL should have indicators"
);
assert!(
result.indicators.len() <= MAX_INDICATORS,
"Should have at most {} indicators, got {}",
MAX_INDICATORS,
result.indicators.len()
);
}
#[test]
fn test_predict_url_detailed_normal() {
let model = crate::load_default_model().expect("Failed to load model");
let result = predict_url_detailed("https://www.google.com", &model);
assert!(
result.score < 0.20,
"Normal URL should score < 0.20, got {}",
result.score
);
assert!(
result.indicators.len() <= MAX_INDICATORS,
"Should have at most {} indicators",
MAX_INDICATORS
);
}
#[test]
fn test_indicator_limit() {
let model = crate::load_default_model().expect("Failed to load model");
let urls = vec![
"http://192.168.1.1/login/verify/account?password=123&token=abc",
"http://very-long-phishing-url-with-many-subdomains.a.b.c.example.com/login",
];
for url in &urls {
let result = predict_url_detailed(url, &model);
assert!(
result.indicators.len() <= MAX_INDICATORS,
"URL '{}' produced {} indicators (max {})",
url,
result.indicators.len(),
MAX_INDICATORS
);
}
}
#[test]
fn test_find_sensitive_keyword() {
assert_eq!(
find_sensitive_keyword("http://example.com/login"),
Some("login".to_string())
);
assert_eq!(
find_sensitive_keyword("http://verify.account.com"),
Some("verify".to_string())
);
assert_eq!(find_sensitive_keyword("http://example.com"), None);
assert_eq!(find_sensitive_keyword("http://bloglogin.com"), None);
assert_eq!(find_sensitive_keyword("http://pineapple.com"), None);
}
#[test]
fn test_ip_url_indicators() {
let model = crate::load_default_model().expect("Failed to load model");
let result = predict_url_detailed("http://192.168.1.1/login", &model);
let has_ip_indicator = result
.indicators
.iter()
.any(|i| i.description.contains("IP address"));
assert!(
has_ip_indicator,
"IP-based URL should have IP address indicator, got: {:?}",
result.indicators
);
}
}